跳到主要内容

数据截至 (上游 commit 32e301ffaf5a)

02 · Engine 装配与配置体系

这一章讲什么： 从 deepspeed.initialize() 的一次调用出发，看 DeepSpeed 怎么把裸模型装配成训练引擎——配置怎么解析、优化器包装链怎么选、训练循环的三个入口（forward/backward/step）各自包了什么。ZeRO 本身的原理在第 1、3 章；本章讲「壳」。

1. 它要解决的小问题​

用户的输入是三样散装货：一个 nn.Module、（可选的）一个优化器、一份 JSON。期望的输出是一个会自己管事的训练器：分布式初始化、混合精度、梯度累积、通信、checkpoint……都不用用户写。

难点在于组合爆炸：fp16/bf16/AMP/fp32 × ZeRO stage 0/1/2/3 × offload 开关 × 用户优化器/配置优化器，每种组合需要不同的包装。DeepSpeed 的答案是：一条由配置驱动的装配链，在 initialize 里一次走完，之后运行时不再分支。

2. initialize：一次装配的完整顺序​

入口 initialize（deepspeed/__init__.py:93）的顺序值得记住，因为所有「为什么我的设置没生效」都能在这条链上定位：

initialize(args, model, optimizer=None, config="ds_config.json", ...)
   │
   ├─ ① 关 zero.Init 上下文（若开着）＋ 断言 model 非空
   ├─ ② dist.init_distributed()          初始化通信后端（comm/comm.py:792）
   ├─ ③ 解析 config 三个来源互斥：args.deepspeed_config / config= / config_params=
   ├─ ④ config_class = DeepSpeedConfig(config, mpu)   JSON → 配置对象
   └─ ⑤ 按模型类型分流造 engine：
          普通模型        → DeepSpeedEngine        (engine.py:252)
          hybrid engine   → DeepSpeedHybridEngine  （训推一体）
          PipelineModule  → PipelineEngine         (pipe/engine.py:60)
   返回 (engine, optimizer, training_dataloader, lr_scheduler)

分流判断是 isinstance(model, PipelineModule)（deepspeed/__init__.py:213），PipelineEngine 分支在 :241 起。返回值固定四元组，用不到的位是 None。

3. 配置体系：一份 JSON 怎么变成对象​

3.1 两层结构​

DeepSpeedConfig（deepspeed/runtime/config.py:692）接受 dict、hjson 文件路径、甚至 base64 字符串，之后 _initialize_params（:796）把字典摊成几十个属性。它是手写 getter 与 pydantic 子配置的混合体：

  • 大子系统有独立的 pydantic 模型：zero_config → DeepSpeedZeroConfig（deepspeed/runtime/zero/config.py:90）、float16_config、bfloat16_config、monitor_config、activation_checkpointing_config 等。
  • 散装键用 get_scalar_param(param_dict, KEY, DEFAULT) 逐个抠出。

引擎侧不直接碰配置字典，而是包一层同名方法：engine.zero_optimization_stage()（deepspeed/runtime/engine.py:1343）只是 self._config.zero_optimization_stage 的转发。全引擎几百个这种转发方法，读代码时认准「engine 方法 = config 属性的透传」即可。

3.2 batch 三件套的反推​

配置里 train_batch_size、train_micro_batch_size_per_gpu、gradient_accumulation_steps 只需给两个，第三个由 _set_batch_related_parameters（config.py:942）按恒等式反推：

train_batch_size = train_micro_batch_size_per_gpu × gradient_accumulation_steps × world_size

三个都给就直接用，一个都不给会撞上 _do_error_check（:1017）的断言。算不清 batch 时先查这三件套和 world_size。

3.3 配置的冲突检查都在初始化早期​

_do_error_check（config.py:1017）和 _do_optimizer_sanity_check（engine.py:1980）把非法组合挡在训练开始前：

  • bfloat16 和 fp16 不能同开（config.py:828 附近的 assert）。
  • AMP 与 ZeRO 互斥（engine.py:1985-1987）。
  • ZeRO + 未测试过的优化器需要显式 zero_allow_untested_optimizer: true（engine.py:1990-1992）。
  • ZeRO-Offload 配了用户自己的优化器时，默认强制换成 DeepSpeedCPUAdam，否则报错并给出明确指引（engine.py:2043-2048）——因为梯度分片是 CPU 张量，普通 GPU 优化器跑不动也跑不快。

4. 优化器包装链：配置怎么选出最终优化器​

_configure_optimizer（engine.py:2028）是装配的心脏。它分两步：

第一步，拿到「基础优化器」。 用户没传优化器就按配置的 optimizer.type 造（_configure_basic_optimizer）；传了就直接用。注意此时它还是裸的，只看全量参数。

第二步，按 _do_optimizer_sanity_check（:1980）的判决结果包一层。 判决逻辑是配置的组合函数：

判决触发条件最终优化器
ZERO_OPTIMIZATIONzero_optimization.stage ≥ 1stage 1/2 → DeepSpeedZeroOptimizer；stage 3 → DeepSpeedZeroOptimizer_Stage3（engine.py:2366 分流）
AMP开了 apex AMPamp.initialize 包装
FP16 / DDP_BFLOAT16模型 fp16/bf16，梯度同精度累积FP16_Optimizer（deepspeed/runtime/fp16/fused_optimizer.py:33，管 loss scale）
BFLOAT16bf16 模型 + fp32 梯度累积BF16_Optimizer（deepspeed/runtime/bf16_optimizer.py:37）
None纯 fp32裸优化器直接用

一个容易踩的分叉：bf16 + ZeRO stage 1 + fp32 梯度累积会被判成 BFLOAT16 而非 ZERO_OPTIMIZATION（engine.py:1995-1998）——即走 BF16_Optimizer 而非 ZeRO 优化器。看到 stage 1 没生效时先查这个组合。

5. 引擎本体：init 装配顺序​

DeepSpeedEngine.__init__（engine.py:255）很长，但顺序固定、每段职责单一：

__init__(args, model, optimizer, ..., config, config_class)
   │
   ├─ _do_args_sanity_check / _configure_with_arguments   （args + mpu 归位，:1630）
   ├─ _configure_expert_parallel                           （MoE 组，:570）
   ├─ _set_distributed_vars                                （global_rank / world_size，:1617）
   ├─ MonitorMaster                                        （监控，可选）
   ├─ _configure_distributed_model(model)                  （★ 模型归位，:1751）
   │     ├─ 半精度 cast（fp16/bf16；ZeRO-3 + zero.Init 的模型跳过）
   │     ├─ module.to(device)
   │     ├─ 扫描 MoE 层，登记 gate/expert 模块
   │     └─ _broadcast_model（:1717，非 ZeRO-3 分片参数广播对齐各 rank）
   ├─ deepspeed_io（:2619）                                （给了 training_data 才建 dataloader）
   ├─ _configure_optimizer（:415 调用点 → :2028）          （★ 上一节的包装链）
   ├─ _configure_lr_scheduler
   └─ checkpoint / flops profiler / engine timers ……

两个细节值得单独记住：

  • engine 是 nn.Module，用户模型挂在 engine.module 上，并通过 __getattr__（:974）把未知属性透传给内部模型——所以 engine.forward(x) 和 model(x) 行为一致，用户的 model.some_custom_method() 也还能调。
  • param_names 在这里建立（{param: name} 映射），ZeRO-3 的 fp32 权重重建、checkpoint 都靠它把分片还原成有名字的 state_dict。

6. 训练循环三入口：壳里包了什么​

6.1 forward（engine.py:2808）​

基本是直通 self.module(*inputs, **kwargs)，外加两件小事：清掉上一轮的 backward 标记；给 loss 注册输出 hook（register_output_backward_hooks），让「用户直接 loss.backward() 而不走 engine.backward()」也能触发 DeepSpeed 的 backward 前置动作。ZeRO-3 的参数调度全在 module hook 里（第 3 章），engine 的 forward 自己不碰参数。

6.2 backward（engine.py:3212）​

顺序固定：loss 按梯度累积步数缩放（scale_wrt_gas）→ fp16 时乘 loss scale → 真正的 loss.backward() → _backward_epilogue（:2977）。

epilogue 里做梯度规约的分叉：

  • 传统路径（stage 0/1 且无 ZeRO 梯度分片）：allreduce_gradients() 对全量梯度 all-reduce。
  • ZeRO 路径：真正的 reduce 早在 backward 过程中被梯度 hook 干完了（第 1 章 §5），epilogue 只调优化器的收尾（backward_epilogue/exit_backward）。

_backward_prologue（:2943）还会把「当前是否梯度累积边界」传给优化器——ZeRO-2/3 在非边界的 micro-step 里也照常 reduce（通信摊平），只是不 step。

6.3 step（engine.py:3412）与梯度累积边界​

step 的核心判断是 is_gradient_accumulation_boundary()（:3270）：当前 micro-step 是否是累积的最后一步。两种模式：

  • managed（默认，managed_gradient_accumulation: true）：引擎内部数 micro_steps，每到 gradient_accumulation_steps 的倍数就是边界。边界才 _take_model_step（:3333）→ optimizer.step() → lr_scheduler.step()；非边界只累计梯度。
  • unmanaged（managed_gradient_accumulation: false）：backward() 只做本地累积，step() 本身就是边界——每次调用都先 finalize 梯度再更新。这给想自己控制累积节奏的训练循环（典型如 RL 框架）留了口。

计数器分三层：micro_steps（每个 micro-batch +1）、global_steps（每次真正更新 +1）、global_samples（按样本数累计）。吞吐计时 tput_timer 挂在 step 尾部。

7. 关键细节与坑​

  • 「我的配置没生效」九成是装配顺序问题。 配置在 initialize 时一次性消费完；运行时改 JSON 字典不会影响引擎。动态能改的只有少数运行参数（lr 走 scheduler，batch 相关有 set_train_batch_size 等）。
  • engine 与 module 的 .to() / .half() 不要混用。 dtype 转换在 _configure_distributed_model 里已经按配置做完；用户再手动 .half() 会把 fp32 buffer（如 rotary 的 inv_freq）也转掉，引擎是刻意只转参数的（:1757 附近注释）。
  • 优化器没传 = 配置里必须有。 两者都没有时 ZeRO 路径会用 DummyOptim 占位（engine.py:2385），只能跑参数调度不能更新——推理/评估场景可用，训练会在 step 的断言处炸。
  • wall_clock_breakdown: true 会开细粒度计时器，定位「时间花在哪」时第一优先开它，但本身就是开销（EngineTimers，:198）。
  • stage 1/2 与 stage 3 的优化器类完全不同，行为差异（如 optimizer.param_groups 里装的是分片还是占位）不要跨 stage 臆测；判断当前是哪条路线，看 type(engine.optimizer).__name__。

8. 代码地图​

主题文件路径符号名
总入口deepspeed/__init__.py:93initialize
分布式初始化deepspeed/comm/comm.py:792init_distributed
配置对象deepspeed/runtime/config.py:692DeepSpeedConfig、_initialize_params、_set_batch_related_parameters
ZeRO 子配置deepspeed/runtime/zero/config.py:90DeepSpeedZeroConfig
引擎类deepspeed/runtime/engine.py:334DeepSpeedEngine
模型归位deepspeed/runtime/engine.py:1835_configure_distributed_model、_broadcast_model
优化器判决deepspeed/runtime/engine.py:2064_do_optimizer_sanity_check
优化器装配deepspeed/runtime/engine.py:2112、:2366_configure_optimizer、_configure_zero_optimizer
三入口deepspeed/runtime/engine.py:2892、:3212、:3412forward、backward、step
累积边界deepspeed/runtime/engine.py:3373、:3333is_gradient_accumulation_boundary、_take_model_step
属性透传deepspeed/runtime/engine.py:1058__getattr__
dataloaderdeepspeed/runtime/engine.py:2703deepspeed_io