DeepSpeed在智星云 GPU 上实战教程完整版

他回疯人院了

278人浏览 · 2026-04-03 11:38:21

他回疯人院了 · 2026-04-03 11:38:21 发布

本教程整理了可直接在智星云平台复制运行的全套代码、配置文件、启动命令和避坑指南，覆盖 ZeRO 全阶段、单卡 Offload、MoE 训练、BERT 微调、断点续训，完全适配智星云预装 CUDA/Anaconda 环境。

0. 前置：智星云实例快速初始化（必跑）

登录实例后，一键配置 DeepSpeed 环境（适配平台预装驱动）

# 创建虚拟环境（推荐）
conda create -n deepspeed python=3.9
conda activate deepspeed

# 安装依赖（智星云 CUDA 11.x 通用）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install deepspeed transformers datasets accelerate

# 验证安装
deepspeed --version
nvidia-smi

1. 核心配置文件（直接复制使用）

1.1 ZeRO-2 配置（多卡训练，性价比最高）

ds_zero2_config.json

{
  "train_batch_size": 32,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "weight_decay": 0.01
    }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "none"
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 1e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 1e8
  },
  "fp16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "checkpoint": {
    "use_node_local_storage": true,
    "tag": "global_step"
  }
}

1.2 ZeRO-Offload 单卡配置（单卡训大模型）

ds_offload_config.json

{
  "train_batch_size": 8,
  "train_micro_batch_size_per_gpu": 2,
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu"
    }
  },
  "fp16": {
    "enabled": true
  },
  "optimizer": {
    "type": "DeepSpeedCPUAdam"
  }
}

1.3 ZeRO-3 全参数分片（超大模型）

ds_zero3_config.json

{
  "zero_optimization": {
    "stage": 3,
    "offload_param": {
      "device": "cpu"
    },
    "offload_optimizer": {
      "device": "cpu"
    }
  },
  "fp16": {
    "enabled": true
  }
}

2. 训练代码模板（通用版，适配所有模型）

train.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification
import deepspeed

# ---------------------- 1. 加载模型与数据 ----------------------
model_name = "bert-large-uncased"  # 可替换为 GPT2/Llama
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# ZeRO-3 必须用这个上下文初始化（关键！）
with deepspeed.zero.Init():
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# ---------------------- 2. DeepSpeed 初始化 ----------------------
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="./ds_zero2_config.json"  # 替换为你的配置文件
)

# ---------------------- 3. 模拟训练数据 ----------------------
device = model_engine.device
texts = ["DeepSpeed 在智星云上训练大模型", "ZeRO 优化解决显存不足"]
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True).to(device)
labels = torch.tensor([1, 0]).to(device)

# ---------------------- 4. 训练循环 ----------------------
model_engine.train()
for step in range(100):
    outputs = model_engine(**inputs, labels=labels)
    loss = outputs.loss
    
    # 反向传播
    model_engine.backward(loss)
    model_engine.step()
    
    if step % 10 == 0:
        print(f"Step {step}, Loss: {loss.item():.4f}")

# 保存 checkpoint
model_engine.save_checkpoint("./checkpoint")

3. 智星云启动命令（直接复制）

3.1 多卡 ZeRO 训练（8/4/2 卡通用）

# 后台运行（防止 SSH 断开中断，必用）
tmux new -s deepspeed

# 激活环境
conda activate deepspeed

# 启动训练（N 卡改 num_gpus）
deepspeed --num_gpus=4 train.py

3.2 单卡 ZeRO-Offload

deepspeed --num_gpus=1 --bind_cores_to_rank train.py

3.3 加载 checkpoint 断点续训

deepspeed train.py --load_dir ./checkpoint

4. MoE 混合专家模型训练（进阶）

启动命令（128 专家，4 卡）

deepspeed --num_gpus=4 train.py \
  --num-experts 128 \
  --moe-expert-parallel-size 4 \
  --moe-train-capacity-factor 1.2 \
  --disable-moe-token-dropping

5. BERT 微调 + DeepSpeed Kernel 加速

deepspeed run_squad.py \
  --model_type bert \
  --model_name_or_path bert-large-uncased \
  --deepspeed ds_zero2_config.json \
  --deepspeed_transformer_kernel \
  --ckpt_type HF \
  --train_micro_batch_size_per_gpu 4 \
  --predict_batch_size 4