数据截至 (上游 commit 32e301ffaf5a)
05 · 选讲：流水并行、MoE、推理与其它特性
这一章讲什么： ZeRO 主线（第 1-4 章）之外的三大子系统，各讲清「入口在哪、核心机制是什么、与 ZeRO 什么关系」，最后给一张其它特性的速查表。本章定位是地图——要用哪块时再顺着代码地图下钻。
1. 流水并行：调度表 + 指令集
它要解决的小问题
模型大到单层都得切着放时，可以把层按段分给不同 rank（流水并行）。难点在调度：micro-batch 在各 stage 间流动，既要填满流水线（少 bubble），又要处理收发依赖。
机制
模型侧：用户把模型写成 PipelineModule（deepspeed/runtime/pipe/module.py:86）——一个层规格的序列，每层用 LayerSpec（:30）声明（类型 + 参数），引擎按 stage 数把层切给各 rank，层只在所属 stage 上真正实例化。
调度侧：PipelineEngine（deepspeed/runtime/pipe/engine.py:60，DeepSpeedEngine 的子类）不跑普通 forward/backward，而是执行一张调度表。训练用 TrainSchedule（deepspeed/runtime/pipe/schedule.py:189）——总步数 2 × (micro_batches + stages - 1)，即经典的 1F1B（one-forward-one-backward）：先连续前向灌满，之后前向反向交替，末尾反向排空。
调度表里的每一步是一串指令对象（PipeInstruction 子类，schedule.py:327-460）：LoadMicroBatch / ForwardPass / BackwardPass / SendActivation / RecvActivation / SendGrad / RecvGrad / ReduceGrads / ReduceTiedGrads / OptimizerStep。执行器是查表分发：_exec_schedule（pipe/engine.py:1396）对每条指令从 _INSTRUCTION_MAP（:1383）找到 _exec_forward_pass（:722）一类实现来跑。stage 间通信走 pipe/p2p.py 的点对点收发。
与 ZeRO 的关系：互斥大部分。 PipelineEngine.__init__ 硬断言 stage < 2（pipe/engine.py:77）——流水并行下参数本来就按 stage 切开了，ZeRO 再切参数/梯度的收益与复杂度不匹配。流水线内部的数据并行副本间仍可用 stage 1 + 梯度 all-reduce（ReduceGrads 指令就是干这个的）。
2. MoE：门控 + all-to-all
它要解决的小问题
MoE（混合专家）把 FFN 换成「N 个专家 + 一个门控」，每个 token 只进 top-k 个专家——参数量暴涨而每 token 计算量不涨。专家并行（EP）把不同专家放不同 rank，于是需要一个「按门控结果把 token 快递到对应专家所在 rank」的通信层。
机制
用户面对的是一个 模块：MoE（deepspeed/moe/layer.py:17），组装三个件——TopKGate（moe/sharded_moe.py:528，按 logits 选 top-k 并做容量裁剪）、Experts（moe/experts.py:13，本 rank 的本地专家副本）、MOELayer（sharded_moe.py:618，调度主体）。
MOELayer.forward（:668）实现 GShard 论文 Algorithm 2，五步：
① 门控：gate(tokens) → 每个 token 的专家编号 + 门控权重 + 辅助 loss
② 编码：按 (专家, capacity) 槽位把 token 摆进 [E, C, M] 的派发张量
③ 第一次 all-to-all：token 跨 rank 送到专家所在处（_AllToAll, :97）
④ 本地专家计算
⑤ 第二次 all-to-all 送回 → 按门控权重加权合并
两个实现细节值得注意：
- 容量（capacity）截断：
capacity_factor决定每专家最多收多少 token，超出的 token 被丢弃（drop_tokens）；C在:686被一次性int(C)物化，注释说明是为了避免后续每次使用都触发一次 device-to-host 同步。 - 与 TP 的组合去重：非专家部分开了张量并行时，各 TP rank 上的 token 是重复的，all-to-all 前先
drop_tokens去重（:705-713注释）——既保证正确性又砍通信量。
与 ZeRO 的关系：共生但分组。 引擎在 _configure_expert_parallel（deepspeed/runtime/engine.py:652）里为 MoE 建独立的专家并行/专家数据并行组；ZeRO 优化器里专家参数与非专家参数走不同的 DP 组（stage_1_and_2.py:778 的 _configure_moe_settings，梯度桶也区分 has_moe_params）。
3. 推理：policy 驱动的内核注入
它要解决的小问题
HuggingFace 的 transformer 层是通用实现：LayerNorm、attention、残差各自独立 kernel，小算子满天飞。推理场景（尤其小 batch）被 kernel launch 与内存读写支配——换融合 kernel 比换硬件便宜。
机制
入口 InferenceEngine（deepspeed/inference/engine.py:43），由 deepspeed.init_inference 造出，核心是注入：
InferenceEngine(model, config)
│ config.injection_policy 或自动识别架构
▼
replace_transformer_layer（module_inject/replace_module.py:189）
│ 按 policy 类（TransformerPolicy, policy.py:43）遍历模型：
│ policy.attention() / policy.mlp() 声明「这层的 qkv/ffn 权重在哪」
▼
换成 containers/ 下对应架构的 DS 容器（bert.py / gpt2.py / bloom.py …）
│ 容器里装的是融合 kernel（attention、LayerNorm、GELU…，csrc/ 里的 CUDA 实现）
▼
可选：TP 切分权重（tensor_parallel tp_size）、量化、CUDA graph
两点取舍：
- policy 是适配层：每种模型架构一个 policy 子类，声明权重怎么从原始模块里抠出来——所以支持的架构是枚举的（
containers/目录即清单），不在清单的架构得自己写 policy 或走auto_tp。 - AutoTP 是另一条路（
module_inject/auto_tp.py:37ReplaceWithTensorSlicing）：不做内核注入，只做张量并行切分——训练侧也能用（第 2 章 engine 的_configure_tensor_parallel）。
本 commit 里还有重写的 inference v2（deepspeed/inference/v2/：engine_v2.py、ragged/ 连续批处理、model_implementations/），与 v1 并存；两条线都在演进，选型以官方文档为准。
4. 其它特性速查
| 特性 | 一句话 | 入口 |
|---|---|---|
| 激活重计算 | 前向不存激活、反向重算，显存换算力 | deepspeed/runtime/activation_checkpointing/checkpointing.py:1058 checkpoint |
| Ulysses 序列并行 | 超长序列沿 sequence 维切，all-to-all 换注意力头 | deepspeed/runtime/sequence_parallel/ulysses_sp.py |
| ZeRO++ | 参数分层分片（zero_hpz_partition_size）+ 量化 all-gather（zero_quantized_weights），砍 ZeRO-3 通信 | deepspeed/runtime/zero/config.py:309-318 |
| MiCS | 亚马逊的分片变体（层次化分片） | deepspeed/runtime/zero/mics.py |
| Muon 优化器 | 矩阵参数的二阶优化器支持 | deepspeed/runtime/zero/muon/ |
| DeepCompile | 把分布式训 练图交给编译器做通信/计算重排 | deepspeed/compile/ |
| 自动调参 | 自动搜 micro-batch / ZeRO 配置 | deepspeed/autotuning/ |
| 弹性训练 | 节点增减时重算 batch 配置 | deepspeed/elasticity/ |
| 监控 | TensorBoard/W&B/COMET 事件写入 | deepspeed/monitor/（MonitorMaster） |
| 通用 checkpoint | ZeRO 分片 checkpoint 重分片加载 | deepspeed/checkpoint/ |
5. 代码地图
| 主题 | 文件路径 | 符号名 |
|---|---|---|
| 流水模型声明 | deepspeed/runtime/pipe/module.py:86、:30 | PipelineModule、LayerSpec |
| 流水引擎 | deepspeed/runtime/pipe/engine.py:60 | PipelineEngine、train_batch、_exec_schedule、_INSTRUCTION_MAP |
| 1F1B 调度表 | deepspeed/runtime/pipe/schedule.py:189 | TrainSchedule、PipeInstruction 一族 |
| MoE 用户入口 | deepspeed/moe/layer.py:17 | MoE |
| 门控与调度 | deepspeed/moe/sharded_moe.py:528、:618、:97 | TopKGate、MOELayer、_AllToAll |
| 引擎的 MoE 登记 | deepspeed/runtime/engine.py:652 | _configure_expert_parallel |
| 推理引擎 | deepspeed/inference/engine.py:43 | InferenceEngine |
| 内核注入 | deepspeed/module_inject/replace_module.py:189、policy.py:43 | replace_transformer_layer、TransformerPolicy |
| AutoTP | deepspeed/module_inject/auto_tp.py:37 | ReplaceWithTensorSlicing |
| 推理 v2 | deepspeed/inference/v2/engine_v2.py | （目录级，自行下钻） |