核心概念：DeepSpeedEngine 与初始化流程

4.1 三个核心抽象

无论你怎么用 DeepSpeed，最后都绕不开这三个对象：

本章把这三个东西讲透，下一章再深入 ZeRO 数学细节。

4.2 deepspeed.initialize() 的签名

从 deepspeed/__init__.py 看，签名是这样的（精简）：

def initialize(args=None,
               model: torch.nn.Module = None,
               optimizer=None,
               model_parameters=None,
               training_data=None,
               lr_scheduler=None,
               distributed_port=29500,
               mpu=None,
               dist_init_required=None,
               collate_fn=None,
               config=None,
               mesh_param=None,
               config_params=None):

实际项目里用得最多的几个参数：

4.3 调用 initialize() 后内部发生了什么

关键点：

• Process group 初始化由 DS 接管，不需要你自己写 torch.distributed.init_process_group；
• JSON 解析失败会立刻抛 pydantic.ValidationError，提示哪一行字段类型不对；
• 注册的 hook 是 DeepSpeed 接管反向、梯度同步、参数搬运的入口。

4.4 DeepSpeedEngine 上你需要知道的方法

4.5 DeepSpeedConfig：JSON → Python 的桥

runtime/config.py 用 pydantic 把 JSON 各字段映射成 Python class。它在做几件事：

1. 类型校验：stage 必须是 0/1/2/3，bf16.enabled 必须是 bool；
2. 互斥检查：fp16.enabled 和 bf16.enabled 不能同时 true；
3. "auto" 解析：把 "auto" 替换成调用方提供的值（HF Trainer 用这个机制注入 batch size 等）；
4. 互相约束：zero_optimization.stage=3 + offload_optimizer.device=nvme 时检查 aio op 是否编译好。

4.6 启动器：deepspeed 命令做了什么

当你写 deepspeed --num_gpus 4 train.py ... 时，deepspeed/launcher/runner.py 会做：

常见 launcher 参数：

deepspeed \
    --num_gpus 8 \                  # 单机用多少卡
    --num_nodes 2 \                 # 多机用几台
    --hostfile hostfile \           # 节点列表
    --master_port 29500 \           # 主节点 NCCL 端口
    --include "worker-0:0,1,2,3" \  # 只用这台机的 0-3 卡
    --no_local_rank \               # 不向子进程传 --local_rank
    train.py --deepspeed --deepspeed_config ds_config.json

多机训练的难点几乎都在 --hostfile、ssh 免密、防火墙——脚本本身不用动。

4.7 三种 rank：local / global / data parallel

DeepSpeed 比 PyTorch 多了一些概念：

当 ZeRO 单独使用时 dp_world_size == WORLD_SIZE；当和 Megatron 模型并行混用时（详见第 9 章）就不等。

4.8 一份等效"PyTorch-style 心智模型"

新手刚切到 DeepSpeed 经常困惑"我以前那套写法去哪了"。可以这样建立映射：

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

• 知道 DeepSpeedEngine = "你模型的封装 + ZeRO 优化器 + 通信 hook"；
• 能从 engine.module 拿到原模型；
• 知道 backward / step 一定要用 engine 的方法；
• 会区分 local_rank、global_rank、dp_world_size；
• 看 launcher 命令时知道每个 flag 的角色。

4.10 DeepSpeedEngine.__init__ 的真实构造顺序

翻开 runtime/engine.py，DeepSpeedEngine 的类定义在第 208 行附近，构造函数大约 100 行，内部把工作切成了 5 个 _configure_* 阶段，调用顺序是固定的（弄清这个顺序，调试 NaN / OOM / hang 时就知道在哪里下断点）：

顺序意义：BF16 wrapper 在 ZeRO wrapper 外面。loss scale / dtype cast 发生在 ZeRO 切分之前看到的张量上。如果你看到一段 stack trace 里出现 BF16_Optimizer.step → DeepSpeedZeroOptimizer.step 的嵌套，就是这个顺序在体现。

4.11 Engine 注册的 hook 总表

DeepSpeed 在初始化时给原 model 挂了若干 hook，这是它能"无感"接管 forward / backward / save 的核心。手写训练循环时调错方法（例如直接 loss.backward()）就会跳过这些 hook，导致 ZeRO-3 hang 或梯度对不齐。

4.12 backward() 的真实路径

当你写 engine.backward(loss) 时，engine.py 第 2615 行附近的代码做了这些（按时间顺序）：

1. 置位 self._running_engine_backward = True（标记自己进入"DS 内部反向"，避免一些 hook 重入）；
2. 调 self.optimizer.scale_if_loss(loss)（ZeRO 路径上会做 loss scale，BF16 时不缩；FP16 时按当前 dynamic loss scale 乘进去）；
3. 真正调 loss.backward(**backward_kwargs)；触发的 backward 路径上，每一层的 hook 是 4.11 注册的那些；
4. 调 self._backward_epilogue()：刷 ZeRO bucket、reduce-scatter 梯度、释放参数（stage 3）；
5. 返回 gas_scaled_loss（按梯度累积步数除回去后的 loss，方便你打日志）。

注意：_scale_wrt_gas 标记控制要不要再按 gradient_accumulation_steps 除一次。HF Trainer 会传 scale_wrt_gas=False（因为 Trainer 自己已经按 GAS 缩了 loss），自己写循环时默认是 True。两边都缩就是常见的"loss 莫名变小到 1/4"bug 来源。

4.13 梯度累积边界判定

DeepSpeed 用一个看起来奇怪、其实很合理的公式判定 "这次 step 是不是真要更新参数"（engine.py:~2679）：

self._is_gradient_accumulation_boundary = (
    (self.micro_steps + 1) % self.gradient_accumulation_steps() == 0
)

关键点：

• 是 micro_steps + 1 而不是 micro_steps：因为 micro_steps 在本次 backward 之后才会自增；要在 backward 之前看到"这次累完了没有"，必须 +1；
• "非边界"的 step() 调用只做梯度累积，不真正更新参数也不更新 lr_scheduler；
• engine._step_applied 反映上一次 step() 是否真的更新过参数——日志里需要"只在真更新的步打 lr"时，用它判 boolean。

4.14 RLHF 场景下的 DeepSpeedHybridEngine

读 deepspeed/__init__.py 第 202 行附近会看到一个分支：

if config.hybrid_engine.enabled:
    engine = DeepSpeedHybridEngine(args=args, model=model, ...)
else:
    engine = DeepSpeedEngine(args=args, model=model, ...)

DeepSpeedHybridEngine 是给 DeepSpeed-Chat 这类 RLHF 训练用的——它内部维护两套权重视图：

• 训练视图：ZeRO 切分态，按 stage 反向；
• 生成视图：合并 + fused inference kernel，actor 做 rollout 时切到这里跑得快。

切换由 engine.eval()/train() 触发。开关字段是 ds_config.json 里的 "hybrid_engine": { "enabled": true, "inference_tp_size": ..., "release_inference_cache": ... }。如果你不做 RLHF，不需要关心这条分支；但看 issue 时偶尔会看到"Hybrid Engine"这个词，知道是这个意思就够了。