Pipeline 引擎 / 3D 并行 / MoE

9.1 为什么除了 ZeRO 还需要别的并行

ZeRO 是"数据并行 + 切优化器/梯度/参数"。它的瓶颈是：

• 不切计算图：每张卡仍然完整 forward 一份模型；万亿参数光 forward 内的激活就装不下；
• 所有卡都参与所有层：跨机带宽不够时，all-gather 时间占比飞涨；
• MoE 专家路由不是 ZeRO 能直接表达的并行。

所以工业训练在 ZeRO 之上还要叠：

9.2 Pipeline 引擎：把模型切成 stage

Pipeline Parallel 的核心：把模型的层数 $L$ 分成 $S$ 段，每段放在一组 GPU 上。前向时数据像水管一样流过 stage 0 → 1 → ... → S-1。

关键效率指标——pipeline bubble 比例：

$$\mathrm{bubble} = \frac{S - 1}{m + S - 1},\qquad m = \text{每步 micro batch 数}$$

所以要让 PP 高效，必须把 micro batch 数 $m$ 设得远大于 stage 数 $S$。

怎么写 PP 训练脚本

DS 的 Pipeline 引擎需要你用 PipelineModule 重新组织模型：

from deepspeed.pipe import PipelineModule, LayerSpec

# 把每层写成一个 LayerSpec（延迟实例化，节省显存）
specs = [
    LayerSpec(EmbeddingLayer, vocab_size, hidden_size),
    *[LayerSpec(TransformerLayer, ...) for _ in range(num_layers)],
    LayerSpec(LMHeadLayer, hidden_size, vocab_size),
]

model = PipelineModule(
    layers=specs,
    num_stages=8,                         # 分成 8 段
    partition_method="uniform",           # 均匀分；其他：parameters / type:Transformer
    loss_fn=lambda outputs, labels: ...,
    activation_checkpoint_interval=1,
)

engine, _, _, _ = deepspeed.initialize(
    args=args, model=model, model_parameters=model.parameters()
)

# 训练循环用 train_batch() 而不是手写 forward/backward
loss = engine.train_batch(data_iter)

注意：

• 不再调 backward / step，而是 train_batch()，它会一次性塞 $m$ 个 micro batch；
• 数据迭代器要在所有 stage 上都拿到对应 micro batch；
• loss 计算只在最后一个 stage 进行。

9.3 3D 并行：ZeRO × TP × PP

千亿模型训练几乎都是 3D 并行：

关系：

$$W = \mathrm{TP} \times \mathrm{PP} \times \mathrm{DP}$$

DS 自己只完整实现 PP + ZeRO；TP 通常借 Megatron 的实现。HF 内部的 Megatron-DeepSpeed fork 就是把两者拼起来训了 MT-NLG 530B。

DS 自带的 tensor parallel

deepspeed/runtime/tensor_parallel/ 是 DS 近几年的新模块。比起 Megatron-LM 的 TP 更轻量，能配合 ZeRO 用。配置举例：

{
  "tensor_parallel": {
    "tp_size": 4,
    "tp_grain_size": 8
  },
  "zero_optimization": {
    "stage": 2
  }
}

仓库 examples/deepspeed/ds_z2_autotp_config.json（如 LLaMA-Factory 也带了一份）就是 ZeRO-2 + 自动 TP 的实验性配置。

9.4 MoE：专家并行

MoE（Mixture of Experts）每层有 $E$ 个 expert，每个 token 路由到 top-$k$ 个 expert。问题：

• $E$ 通常很大（8/16/64/256），不可能每张卡都存所有 expert；
• 需要把 expert 切到不同卡，token 跨卡通信完成。

DS 在 deepspeed/moe/ 提供 MoE layer：

from deepspeed.moe.layer import MoE

class MoEBlock(nn.Module):
    def __init__(self, dim, num_experts, ep_size):
        super().__init__()
        self.moe = MoE(
            hidden_size=dim,
            expert=Expert(dim),
            num_experts=num_experts,
            ep_size=ep_size,          # expert 并行度
            k=2,                       # top-2 路由
            min_capacity=4,
            use_residual=False,
        )
    def forward(self, x):
        out, l_aux, exp_counts = self.moe(x)
        return out, l_aux

l_aux 是负载均衡辅助 loss，要加到主 loss 上：

loss = main_loss + 0.01 * l_aux

在 ds_config 中相关字段：

{
  "zero_optimization": {
    "stage": 2,
    "moe_layer_size_threshold": 0,
    "moe_min_capacity": 4
  }
}

9.5 Sequence Parallel：长上下文必备

当 sequence length $L$ 很长（32K / 128K / 1M），单卡装不下 attention 矩阵（$O(L^2)$）。SP 把序列沿 seq 维度切到不同卡：

DS 中开启：

{
  "sequence_parallel_size": 4
}

9.6 实战层次推荐

9.7 与 Megatron-LM 的关系

第 1 章已经讲过 DS 和 Megatron 的对比。这里给一张实战上的"分工图"：

它们通过 mpu（model parallel utility） 互相握手——你在 deepspeed.initialize(..., mpu=mpu) 时传入 Megatron 的 mpu 对象，DS 就能知道哪几张卡属于同一个 TP/PP 组，从而正确处理 ZeRO。

9.9 PipelineModule 的构造细节

9.2 提到 PP 训练要"用 PipelineModule 重新组织模型"。这一节展开看里面长什么样。源码在 deepspeed/runtime/pipe/module.py。

9.9.1 LayerSpec（延迟实例化）

# pipe/module.py:30–75
class LayerSpec:
    def __init__(self, typename, *module_args, **module_kwargs):
        self.typename     = typename
        self.module_args  = module_args
        self.module_kwargs= module_kwargs
    def build(self):
        return self.typename(*self.module_args, **self.module_kwargs)

为什么要延迟？因为 PipelineModule 是 N 张卡都跑同一份代码，每张只持有自己 stage 的层；如果在所有 rank 都先把全模型实例化一遍再扔掉别人那部分，峰值显存会等于完整模型，ZeRO-3 的省显存意义就没了。LayerSpec 让每张卡只 build() 自己负责的层。

9.9.2 TiedLayerSpec（权重 tying）

# pipe/module.py:77–84
class TiedLayerSpec(LayerSpec):
    def __init__(self, key, typename, *args, tied_weight_attr=['weight'], forward_fn=None, **kwargs):
        ...

典型场景：embedding 和 LM head 共享同一份权重，但它们位于 PP 的第一个 stage 和最后一个 stage——跨 stage tied。key 是 tying 的唯一标识，tied_weight_attr 列出要 tied 的 attr（默认 ['weight']，bias 一般不 tied）。引擎会在 backward 时跨 stage 把 tied 梯度 all-reduce（对应 9.10 里的 ReduceTiedGrads 指令）。

9.9.3 partition_method 三种取值

9.10 Pipeline schedule 全表 & PipeInstruction

9.2 那张 1F1B 示意图实际对应 deepspeed/runtime/pipe/schedule.py 的 TrainSchedule（line 189–299）。读这个文件能让 PP 的"为什么各 stage 做的事不一样"完全可视化。

9.10.1 三种 schedule

9.10.2 10 种 PipeInstruction（schedule.py:327–495）

PipelineEngine._exec_schedule()（engine.py:1364 附近）就是一个 for cmd in schedule: _INSTRUCTION_MAP[type(cmd)](cmd) 的 dispatch 循环。看 hang 时可以打印这个循环卡在哪条指令。

9.11 ProcessTopology：3D 世界的笛卡尔坐标

当 DP × TP × PP × EP 同时存在时，"我是哪张卡"已经不能用一个全局 rank 说清楚——你属于某个 DP 组的第几个、某个 PP 段的第几个、某个 TP 内的第几个。deepspeed/runtime/pipe/topology.py（line 12–126）的 ProcessTopology 就是这套多维坐标系。

# 形如：3D 训练 24 张卡，PP=4 / DP=3 / TP=2
topo = ProcessTopology(
    axes=['pipe', 'data', 'model'],
    dims=[4, 3, 2]
)
# 拿到 (pipe=2, data=1, model=0) 对应的全局 rank
g_rank = topo.get_rank(pipe=2, data=1, model=0)
# 沿 data 轴的所有 rank（用来建 DP 通信组）
dp_ranks = topo.get_axis_comm_lists('data')

关键方法：

PipelineParallelGrid 是 ProcessTopology 的训练侧包装，提供 get_stage_id() / get_data_parallel_group() / get_pipe_parallel_group() 等便捷接口。多机调度的 P2P 通信靠 deepspeed/runtime/pipe/p2p.py，里面 send(tensor, dest_stage) / recv(tensor, src_stage) 是真正发收的两个函数。

9.12 MoE 关键参数源码解读

9.4 给出了 MoE layer 的用法，这里展开三个最容易踩坑的参数。源码在 deepspeed/moe/sharded_moe.py 与 deepspeed/moe/layer.py。

9.12.1 capacity_factor

每个 expert 在一个 batch 内只接收固定数量的 token（avoid OOM）：

# sharded_moe.py:168
capacity = ceil((num_tokens / num_experts) * capacity_factor)

capacity_factor=1.0 表示理想均匀分布——但路由器很难做到 100% 均匀，所以默认 1.0 一定会丢 token。常见取值 1.25–2.0。drop_tokens=False 时（line 216）capacity 自动扩到 max(exp_counts) 不丢任何 token，代价是显存峰值由最忙的 expert 决定。

9.12.2 Random Token Selection（RTS）

当某 expert 路由数 > capacity 时，谁被丢？传统做法是按 router 分数排，分数低的丢。RTS（sharded_moe.py:233）改成随机选该 expert 容量内保留的 token，理由：避免路由器固化"偏爱某些位置"，提高负载均衡。打开 use_rts=True 即用。

9.12.3 l_aux 辅助损失的公式

# sharded_moe.py:231
l_aux = sum_e( mean(gates[e]) * mean(mask1[e]) ) * num_experts

mean(gates[e]) 是 expert e 收到的平均路由概率，mean(mask1[e]) 是 expert e 实际被选 top-1 的频率。两者乘积反映"概率分布与实际选择是否一致"，乘以 num_experts 归一到 [1, +∞)。训练时必须手动加进总 loss：

loss = main_loss + 0.01 * l_aux   # 0.01 是经验权重，论文常用范围 0.01–0.1

不加的话路由器没动力均衡，几步之后所有 token 都路由到同一个 expert，其余 expert 死掉。

9.12.4 _AllToAll autograd Function

MoE 的核心通信是 token 在 expert 之间跨卡重新分发，用的是自定义的 _AllToAll（sharded_moe.py:97–109）：

class _AllToAll(torch.autograd.Function):
    @staticmethod
    def forward(ctx, group, input): ...   # dist.all_to_all_single
    @staticmethod
    def backward(ctx, grad_output): ...   # 反方向再 all_to_all 一次

这个 Function 的存在让 MoE 的 backward 自动对齐 token 路由，不需要用户写任何通信代码。

9.13 AutoTP：怎么自动把 HF 模型切成 TP

2025-03 引入的 AutoTP（deepspeed/module_inject/auto_tp.py）的目标是让用户不写一行 policy 代码就能 TP。核心流程：

QKV fused 切法的麻烦在于：保存的 checkpoint 把 Q/K/V 三个矩阵按 hidden 维拼起来，TP 又要按 head 维切。strided_copy 的作用就是先按外维切 3 份 → 再在每份内部按 head 切 → 最后按 TP rank 重组。读这一段是理解 AutoTP 的难点也是它最聪明的地方。

9.14 Sequence Parallel：Ulysses 与 ALST

9.5 概述了 SP 的两种风格——这里展开 DS 自家 Ulysses 的实现细节。源码在 deepspeed/sequence/。

9.14.1 attention 沿 head 切 + 两次 all-to-all

Ulysses 的核心做法：在 attention 前后各做一次 all_to_all_single，前者把"seq 分布 / head 完整"变成"seq 完整 / head 分布"，后者反过来。这样 attention 计算用的是完整 sequence（normal attention 不变），但每张卡只算一部分 head。

9.14.2 Arctic Long Sequence Training（ALST）

2025-06 上线的 ALST 把 Ulysses 推到百万 token 级别。新加了一个文件 deepspeed/sequence/fpdt_layer.py（约 60 KB），把 FlashAttention 风格的分块 attention、动态稀疏、激活分块下放到 sequence 子系统。同目录的 auto_sp.py 是 "auto sequence parallel"——在 model trace 时识别 attention 模块，自动插入 SP 包装，类似 AutoTP 的思路。

开启入口：ds_config.json 顶层加 "sequence_parallel_size": 4（按你的 SP 度），DS 在 deepspeed.initialize() 时检测到非 1 就调 auto_sp 包模型。

9.15 Domino：把 TP 通信推到 backward 才同步

2024-11 引入的 Domino（deepspeed/runtime/domino/）解决一个具体问题：传统 TP 列并行的 forward 后必须 all-reduce，计算和通信是串行的。Domino 让 all-reduce 异步发出但不立刻 wait，等到 backward 第一次用到这块梯度时再 wait。

挂法：当前要用 Domino 必须改模型代码——把 nn.Linear 换成 DominoAsyncColumnParallelLinear / RowParallelLinearNoComm。DS 暂时不通过 config 一键启用，主要作为 "想要 TP 极致性能 + 不怕改代码" 的进阶选项。

9.16 这章你需要带走什么

• 知道 ZeRO 只是数据并行，真正的 3D 训练还要 TP/PP/EP/SP；
• 知道 DS 的 Pipeline 用 PipelineModule 重写模型，LayerSpec / TiedLayerSpec 是延迟实例化的关键；
• 能在 hang / 性能问题时按 PipeInstruction 表对应到具体源码行；
• 知道 MoE 路由的 l_aux loss 不能忘加，capacity_factor / RTS 怎么取；
• 知道 AutoTP 的注入流程靠的是 tp_parser + replace_policies + strided_copy；
• 知道 Ulysses 是两次 all-to-all + head/seq 维度切换；Domino 是把 TP allreduce 推迟到 backward；
• 知道工业级千亿训练通常走 Megatron-DeepSpeed，而不是单纯 DS。