快速上手：第一次跑通 DeepSpeed

3.1 两条路径

大多数实际项目走 B 路径（最省力）；但理解 A 路径才能在 B 出错时知道发生了什么。本章先讲 A，再讲 B。

3.2 路径 A：原生 DeepSpeed 训练脚本

3.2.1 一份最小可见的训练脚本

# train.py
import torch
import torch.nn as nn
import deepspeed
import argparse

class TinyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )
    def forward(self, x): return self.net(x)

def get_args():
    p = argparse.ArgumentParser()
    p = deepspeed.add_config_arguments(p)        # 让 DS 接管 --deepspeed_config
    p.add_argument("--local_rank", type=int, default=-1)
    return p.parse_args()

args = get_args()
model = TinyMLP()

# ★ 关键调用：注入 DeepSpeed 引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=model.parameters(),
)

for step in range(100):
    x = torch.randn(32, 128, device=model_engine.device)
    y = torch.randint(0, 10, (32,), device=model_engine.device)

    logits = model_engine(x)
    loss = nn.functional.cross_entropy(logits, y)

    model_engine.backward(loss)        # ★ 不要 loss.backward()
    model_engine.step()                # ★ 不要 optimizer.step()/.zero_grad()

    if step % 10 == 0:
        print(f"step={step} loss={loss.item():.4f}")

3.2.2 配套 ds_config.json

{
  "train_micro_batch_size_per_gpu": 32,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 1e-3, "weight_decay": 0.01 }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": { "warmup_min_lr": 0, "warmup_max_lr": 1e-3, "warmup_num_steps": 50 }
  },
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 2
  }
}

3.2.3 启动

启动之前先说一下脚本里那行 deepspeed.add_config_arguments(p) 到底做了什么——它给你的 argparse 自动加了下面这几组 flag（实现见 deepspeed/__init__.py 与 deepspeed/launcher/runner.py）：

所以训练脚本里 argparse 不需要你自己写这些 flag，加完 add_config_arguments() 就有了。--local_rank 是另一个独立的、launcher 注入的 flag，必须自己加（见上一节脚本里的那一行）。

# 单卡
deepspeed --num_gpus 1 train.py --deepspeed --deepspeed_config ds_config.json

# 4 卡单机
deepspeed --num_gpus 4 train.py --deepspeed --deepspeed_config ds_config.json

# 2 节点 × 8 卡多机（用 hostfile）
deepspeed --hostfile hostfile --num_gpus 8 \
    train.py --deepspeed --deepspeed_config ds_config.json

其中 hostfile 长这样：

worker-01 slots=8
worker-02 slots=8

3.3 把这份脚本读懂

第一次写 DeepSpeed 代码，最容易踩的坑就是不知道哪几行变了。下表把原版 PyTorch 和 DeepSpeed 版对照：

3.4 deepspeed.initialize() 的 4 个返回值

从源码 deepspeed/__init__.py（initialize() 的实现大致在第 80–252 行）看，initialize() 返回的是一个固定 4 元组：

model_engine, optimizer, train_dataloader, lr_scheduler = deepspeed.initialize(...)

分别是：

3.5 整个过程的时序图

3.6 路径 B：HF Trainer + DeepSpeed

用 transformers Trainer 时几乎不用碰 DS 调用，只需：

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    bf16=True,
    deepspeed="ds_config.json",          # ★ 关键
)

trainer = Trainer(model=model, args=training_args,
                  train_dataset=ds_train, tokenizer=tok)
trainer.train()

启动：

torchrun --nproc_per_node 4 train.py
# 或
deepspeed --num_gpus 4 train.py

HF Trainer 内部检测到 training_args.deepspeed 不为空时，会自动调 deepspeed.initialize() 包住 model / optimizer / scheduler，train() 内的 loss.backward() 会被替换成 model_engine.backward()。所有"路径 A"做的事都被 Trainer 隐藏了。

这就是为什么第 1 章说 DeepSpeed 在 HF Trainer 用户眼中"看不见但都在用"。LLaMA-Factory（见 那本书 ch10）也是基于这条链路。

3.7 把"auto"用起来

HF Trainer 配 DS 时，ds_config.json 大量字段填 "auto"：

{
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "train_batch_size": "auto",
  "gradient_clipping": "auto",
  "bf16": { "enabled": "auto" },
  "zero_optimization": { "stage": 3 }
}

含义：让 Trainer 用 YAML / TrainingArguments 里的对应字段填上。这样同一份 JSON 给不同任务复用，只在 YAML 里改 batch 即可。这种"配置注入"是 HF Trainer 集成的精髓，第 8 章会拆解。

3.8 这章给你留下的应该是什么

• 知道 deepspeed.initialize() 是入口，返回 4 个对象；
• 知道 .backward() / .step() 必须用 engine 的方法，不能用 PyTorch 原生；
• 知道 deepspeed launcher 怎么拉起多卡 / 多机；
• 会看 ds_config.json 的最简结构（optimizer + scheduler + bf16 + zero_optimization）；
• 暂时还不需要懂 ZeRO stage 各自切了什么——下一章会讲。