分布式 Transformer 深度解析：从张量形状到通信拓扑

前情说明：技术深度 D3。起点是 ailzhang 的可视化博客——她把 DP/TP/SP/CP/EP/VP 六种并行在 Transformer 各 block 里的张量形状和通信原语画得非常清爽，站在单张 GPU 视角看 sharding 的 dance。推荐先读她那篇看图理解直觉，我这篇是补深度：

每种并行的 tensor 形状分步演变 + 通信量在哪一步产生——ASCII 图逐层展开

通信原语底层怎么走——Ring AllReduce、Tree、Halving-Doubling、AllToAll

FSDP / ZeRO 1/2/3 / Pipeline Parallel 这些 ailzhang 明确留给下一篇的方案

从训练侧扩到推理侧——TP/EP/PP 在推理里的差异

目标读者：已经知道 Transformer 基本结构、想搞懂"为什么大家都这么分卡"的工程师。所有图都是 ASCII,配合等宽字体阅读。

一、先建立坐标系：一次 forward 到底有多贵

在讨论"为什么要切"之前，先算明白"不切能不能行"。

以 Llama 3 70B 做锚点：

符号	含义	Llama 3 70B
B	Global batch size	512 sequences
S	Sequence length	8192
D	Hidden dim	8192
F	FFN dim	28672
L	Layers	80
H	Attention heads	64 (GQA: 8 KV heads)
V	Vocab	128k

1.1 显存账（单卡装不下）

参数（bf16）        : 70B × 2 bytes = 140 GB
优化器状态（Adam fp32）: 70B × 12 bytes = 840 GB   (m + v + master weight)
梯度（bf16）        : 70B × 2 bytes = 140 GB
激活值（80 层）     : ≈ 1-2 TB              (无 checkpointing)
─────────────────────────────────────────
合计               : > 2 TB

单张 H100 才 80 GB。分布式不是"要不要"，是"必须切,问题是怎么切"。

1.2 三种约束的取舍

分布式训练所有设计都在平衡三件事：

       ┌─────────────┐
       │    显存     │  装不下 → 切参数/激活
       │  (MEMORY)   │
       └──────┬──────┘
              │
              │ 三角平衡
    ┌─────────┼─────────┐
    │                   │
┌───▼────┐         ┌────▼───┐
│  算力  │◄───────►│ 通信   │
│(COMPUTE│         │(COMM)  │
│ / MFU) │         │        │
└────────┘         └────────┘
 跑得慢 → 分 batch  卡间等 → 拓扑分层

核心原则：通信量大的并行走机内 NVLink（900 GB/s），通信量小的走机间 InfiniBand（100-400 GB/s）。理解这一句话，就理解了并行拓扑设计的一半。

二、通信原语：底层到底怎么走

ailzhang 的文章反复出现 AllReduce、ReduceScatter、AllGather，但没讲底层。这决定了每个 collective 的实际带宽和延迟——是理解一切并行的地基。

2.1 四个基本 collective

假设 4 张卡（GPU 0/1/2/3），每张持有一个 tensor（大小 M bytes）。看这张对照图:

初始状态：
  GPU 0: [A]     GPU 1: [B]     GPU 2: [C]     GPU 3: [D]

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
AllReduce（每卡都要完整 reduce 结果）:
  GPU 0: [A+B+C+D]  GPU 1: [A+B+C+D]  GPU 2: [A+B+C+D]  GPU 3: [A+B+C+D]

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
ReduceScatter（reduce 后每卡只拿一片）:
  假设每个 tensor 已经切成 4 块，A = [a0,a1,a2,a3] 等等
  GPU 0: [a0+b0+c0+d0]
  GPU 1: [a1+b1+c1+d1]
  GPU 2: [a2+b2+c2+d2]
  GPU 3: [a3+b3+c3+d3]

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
AllGather（每卡有一片，聚成完整）:
  假设 GPU i 持有分片 s_i
  GPU 0: [s0, s1, s2, s3]  &#x3C;── 拿到所有
  GPU 1: [s0, s1, s2, s3]
  GPU 2: [s0, s1, s2, s3]
  GPU 3: [s0, s1, s2, s3]

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
AllToAll（每卡切成 N 份,第 i 份发给卡 i）:
  假设 GPU i 有 [i→0, i→1, i→2, i→3]（4 份要发给 4 张卡）
  GPU 0: [0→0]  ─┐
  GPU 1: [1→0]  ─┼──► GPU 0 收到 [0→0, 1→0, 2→0, 3→0]
  GPU 2: [2→0]  ─┤    GPU 1 收到 [0→1, 1→1, 2→1, 3→1]
  GPU 3: [3→0]  ─┘    GPU 2/3 类似

数学恒等式（背下来）：

    AllReduce  ==  ReduceScatter + AllGather
    ────────       ─────────────    ──────────
    每卡完整sum    每卡一片sum      每卡集齐所有

这是 FSDP / ZeRO / SP 能省显存的核心：AllReduce 拆成两步后，中间那一步是"每卡只有一片"——这一片可以直接用于分布式优化器更新，不需要每卡都拿到完整梯度。ZeRO 的每一级都是在这条恒等式的不同"截断"处做文章。

2.2 Ring AllReduce：为什么带宽不受卡数影响

朴素 AllReduce：所有卡把数据发给一张主卡，主卡加起来广播回去。看图:

朴素 AllReduce（差):
    GPU 0  GPU 1  GPU 2  GPU 3         ┌── 主卡瓶颈 ──┐
      │      │      │      │            │             │
      │      ▼      ▼      ▼            ▼             ▼
      └────► GPU 0 ◄────────┘         接收 3M         发送 3M
             (求和)                    带宽爆炸        带宽爆炸
                │
                └──► 广播回 3 张卡

  卡数 N 增大 → 主卡带宽 O(N) 线性下降,64 卡就跪

Ring AllReduce（NVIDIA NCCL 默认算法）的优雅之处：

Step 0: 每卡把自己的 tensor 切成 N 份（N=4）
  GPU 0: [a0][a1][a2][a3]
  GPU 1: [b0][b1][b2][b3]
  GPU 2: [c0][c1][c2][c3]
  GPU 3: [d0][d1][d2][d3]

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Phase 1: ReduceScatter（N-1=3 步）

Step 1:  每卡把"编号 == 自己编号 - 1"的块发给下一张卡
  GPU 0 [a0]───► GPU 1 (收到 a0 + 累加到本地 b0 变成 a0+b0)
  GPU 1 [b1]───► GPU 2 (收到 b1 + 累加到本地 c1 变成 b1+c1)
  GPU 2 [c2]───► GPU 3 (收到 c2 + 累加到本地 d2 变成 c2+d2)
  GPU 3 [d3]───► GPU 0 (收到 d3 + 累加到本地 a3 变成 d3+a3)

Step 2: 每卡把"上一步刚累加过的块"再传下去（滚雪球）
  GPU 0 [d3+a3]───► GPU 1 变成 d3+a3+b3
  GPU 1 [a0+b0]───► GPU 2 变成 a0+b0+c0
  GPU 2 [b1+c1]───► GPU 3 变成 b1+c1+d1
  GPU 3 [c2+d2]───► GPU 0 变成 c2+d2+a2

Step 3: 再传一次,每卡都持有一个完整 sum
  GPU 0: [_ , _ , a2+b2+c2+d2, _ ]  ← 完整 sum 的第 2 块
  GPU 1: [a0+b0+c0+d0, _ , _ , _ ]  ← 完整 sum 的第 0 块
  GPU 2: [_ , b1+c1+d1+a1, _ , _ ]  ← 完整 sum 的第 1 块
  GPU 3: [_ , _ , _ , d3+a3+b3+c3]  ← 完整 sum 的第 3 块

  ↑ ReduceScatter 完成！每卡持有一片完整 sum

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Phase 2: AllGather（再 N-1=3 步）

Step 4-6: 类似的环形传递，只不过这次不再累加，是把每卡的那片
         广播给其他卡。3 步后每卡都持有完整的 4 片。

总通信量：每卡收发 2(N-1)/N × M bytes ≈ 2M——和 N 无关！

关键洞察:
   卡数 N ↑         →  每卡收发量恒定 ≈ 2M
   卡数 N ↑         →  链路带宽利用率 ↑（所有链路都在工作）
   卡数 N ↑         →  步数 O(N)，延迟 ↑

结论: Ring 适合"大消息 + 中等卡数",瓶颈从"带宽"转到"延迟"

局限：延迟 O(N)，1000+ 卡时不再最优。这时切到 Tree 或 Halving-Doubling。

2.3 Tree AllReduce vs Halving-Doubling

Tree AllReduce（二叉树 reduce + broadcast）:

  Reduce 阶段（log N 步）:            Broadcast 阶段:

        [7]  root                      [7] root  (完整 sum)
        / \                             / \
      [3] [4]                        [3] [4]
      / \  / \                       / \  / \
    [1][2][3][4]                   [1][2][3][4]

  延迟: O(log N)  ← 卡多时更好
  带宽: M（一半 traffic）
  代价: 树的每一层链路利用率低（只有一半的卡在参与）

Halving-Doubling（NCCL 默认在大 N 场景）:

Phase 1: 递归 Halving（log N 步）
  Step 1: 相邻两两配对，各交换一半数据并 reduce
    (0↔1) (2↔3) (4↔5) (6↔7)
    每对分别持有 "低一半 sum" 和 "高一半 sum"
  Step 2: 距离 2 的配对，再交换 sum 的一半
    (0↔2) (1↔3) (4↔6) (5↔7)
  Step 3: 距离 4 的配对
    ...

Phase 2: 递归 Doubling（log N 步）
  反过来 broadcast 分片

延迟: O(log N)  ← 和 Tree 一样好
带宽: 2M       ← 和 Ring 一样好
兼具两者优势,是"大 N + 大 msg"的最佳选择

面试锚点：Ring 适合大消息（GBs 级），Tree 适合小消息（KB 级）+ 大规模，Halving-Doubling 兼顾。NCCL 会根据 msg size × N 自动选，无需手动调参。

2.4 AllToAll：MoE 的通信之王

AllToAll 是 MoE 的核心 collective：每张卡把自己的数据切成 N 份，第 i 份发给第 i 张卡。看图:

初始状态（4 张卡,每张卡持有一段 token 序列,已经算好每个 token 要去哪个 expert）:

  GPU 0: tokens [t0..t9]  路由决策: t0→E1, t1→E3, t2→E0, ...
  GPU 1: tokens [t10..t19] 路由决策: ...
  GPU 2: tokens [t20..t29] 路由决策: ...
  GPU 3: tokens [t30..t39] 路由决策: ...

Step 1: 每卡按目标 expert 做 permutation（本地重排 token）
  GPU 0: [→E0: t2,t5] [→E1: t0,t3] [→E2: t7,t9] [→E3: t1,t4,t6,t8]
  GPU 1: [→E0: t11]   [→E1: t14,t18] [→E2: t10,t15,t17] [→E3: t12,t19]
  GPU 2: [→E0: t23,t28] [→E1: t20] [→E2: t24,t27] [→E3: t21,t22,t25]
  GPU 3: [→E0: t32,t38] [→E1: t33] [→E2: t35] [→E3: t31,t34,t36]

Step 2: AllToAll 一次性交换（每对卡之间同时打开链路）
       ┌───────────────────────────────────────────┐
  GPU 0│──┬──►GPU 0[→E0]        (自己的第 0 段留下)│
       │  ├──►GPU 1[→E1]                           │
       │  ├──►GPU 2[→E2]                           │
       │  └──►GPU 3[→E3]                           │
       │                                            │
  GPU 1│──┬──►GPU 0[→E0]                           │
       │  ├──►GPU 1[→E1]                           │
       │  ├──►GPU 2[→E2]                           │
       │  └──►GPU 3[→E3]                           │
       │  ...                                       │
       │  N² 条链路同时活跃！                       │
       └───────────────────────────────────────────┘

Step 3: 交换完成后，每卡收到该卡上 expert 该处理的所有 token
  GPU 0（跑 E0）: 收到 [t2,t5,t11,t23,t28,t32,t38]
  GPU 1（跑 E1）: 收到 [t0,t3,t14,t18,t20,t33]
  GPU 2（跑 E2）: 收到 [t7,t9,t10,t15,t17,t24,t27,t35]
  GPU 3（跑 E3）: 收到 [t1,t4,t6,t8,t12,t19,t21,t22,t25,t31,t34,t36]

  ↑ 注意每卡收到的 token 数不同！这就是 MoE 负载不均衡问题

Step 4: 各 expert 本地跑 FFN
Step 5: 再一次 AllToAll，把结果 token 送回原位（combine）

通信量：每卡收发 M bytes（数据总量,不因 N 变化）。

为什么头疼：

不能用 ring 优化——N² 条链路同时活跃,fabric 全打开
延迟主导——路由必须完成才能算 FFN,不像 AllReduce 可以 overlap 到反向传播
负载不均——每卡收到的 token 数不确定,拖尾严重

优化方案：

Hierarchical AllToAll：先机内做局部 AllToAll → 机间做一次 → 机内再一次
Fused kernel：permutation + AllToAll + GEMM 融合，隐藏延迟
Expert placement：频繁互相路由的 expert 放同一台机器
DualPipe（DeepSeek V3）：把 AllToAll 塞进 pipeline 空闲气泡

三、六种并行 × 通信量对照表

这张表是这篇文章最实用的东西。以 Llama-3 70B、一步 forward+backward、GBS=512 为例：

并行	切什么	关键 collective	单步通信量（每卡）	建议网络层
DP	Batch B	AllReduce(gradients)	2 × params ≈ 280 GB/step	机间 IB（可 overlap）
TP	Weights (heads/F)	AllReduce(activations) × 4/layer	4 × L × B × S × D × 2 ≈ 21 TB/step	机内 NVLink（必须）
SP	Activations (S) 在 element-wise op	AllGather + ReduceScatter	和 TP 同数量级	机内 NVLink（配 TP 用）
CP	Sequence in Attention	Ring send/recv (KV)	~2 × L × B × S × D × 2	机内 NVLink
EP	Experts	AllToAll × 2/layer	~2 × L × B × S × D × 2（激活分数）	机内/机间都行，但机内更好
PP	Layers	Send/Recv(activations)	2 × B × S × D × (PP-1)	机间 IB（延迟不敏感）
VP	Vocab	AllReduce(loss stats)	极小（几 KB）	任意

关键规律（背下来）：

TP/SP 通信量最大（每层 4 次 collective），必须走 NVLink（机内）。跨机 TP = 灾难
PP 通信量最小（只传激活），可以走 机间 IB。但有 pipeline bubble 问题
DP 通信量中等，但可以 overlap 到 backward 里（NCCL 支持异步），实际不成为瓶颈
EP 通信量大且是 AllToAll，是 MoE 的头号问题
VP 通信量小，随便放

这决定了千卡训练的拓扑：

                    3D 并行经典拓扑
    ┌─────────────────────────────────────────────────┐
    │                                                  │
    │   ┌──────────────┐  ┌──────────────┐            │
    │   │ 机器 A       │  │ 机器 B       │            │
    │   │ ┌──┬──┐      │  │ ┌──┬──┐     │            │
    │   │ │G0│G1│      │  │ │G0│G1│     │            │
    │   │ ├──┼──┤ NVL  │  │ ├──┼──┤ NVL │  ← TP 组   │
    │   │ │G2│G3│      │  │ │G2│G3│     │    机内    │
    │   │ ├──┼──┤      │  │ ├──┼──┤     │            │
    │   │ │G4│G5│      │  │ │G4│G5│     │            │
    │   │ ├──┼──┤      │  │ ├──┼──┤     │            │
    │   │ │G6│G7│      │  │ │G6│G7│     │            │
    │   │ └──┴──┘      │  │ └──┴──┘     │            │
    │   └──────┬───────┘  └──────┬──────┘            │
    │          │                  │                    │
    │          └──── IB ──────────┘   ← PP 组          │
    │             (机间)              (跨机)           │
    │                                                  │
    │        ┌────────────────────┐                    │
    │        │  ...多个 PP 组     │   ← DP 组          │
    │        │  AllReduce gradient│    (跨 pod)        │
    │        └────────────────────┘                    │
    └─────────────────────────────────────────────────┘

    典型: 8 卡机内 TP=8, PP 跨机, DP 跨 pod

四、Tensor Parallel：Column + Row 的免费配对

现在开始一种一种把并行讲透。第一个也是最核心的：TP。

4.1 Column Parallel（切输出维度）

对权重 W ∈ [D, F]，按列切成 W = [W_0, W_1, W_2, W_3]（假设 N=4）,每卡持有 W_i ∈ [D, F/4]。

              W ∈ [D, F]                        Y = X · W ∈ [B*S, F]
     ┌──────┬──────┬──────┬──────┐
     │      │      │      │      │
     │ W_0  │ W_1  │ W_2  │ W_3  │           ┌─────┬─────┬─────┬─────┐
  D  │      │      │      │      │    ─►     │ Y_0 │ Y_1 │ Y_2 │ Y_3 │
     │      │      │      │      │           └─────┴─────┴─────┴─────┘
     │      │      │      │      │            B*S       F/4 each
     └──────┴──────┴──────┴──────┘
       F/4   F/4    F/4    F/4

切分方式:
  GPU 0 持有 W_0  →  Y_0 = X · W_0  ∈ [B*S, F/4]
  GPU 1 持有 W_1  →  Y_1 = X · W_1  ∈ [B*S, F/4]
  GPU 2 持有 W_2  →  Y_2 = X · W_2  ∈ [B*S, F/4]
  GPU 3 持有 W_3  →  Y_3 = X · W_3  ∈ [B*S, F/4]

输入要求:  X 必须每卡都有（"复制"状态）
输出状态:  Y 是分片的（每卡持有 F/4 那一片）,不需要通信

核心：输入是复制的、输出是分片的——结果自然分片，正好适合接下一个 Row Parallel。

4.2 Row Parallel（切输入维度）

对权重 W ∈ [F, D]，按行切成 W = [W_0; W_1; W_2; W_3],每卡持有 W_i ∈ [F/4, D]。

      X ∈ [B*S, F]                W ∈ [F, D]
    分片输入                      ┌──────────┐
    ┌─────┬─────┬─────┬─────┐    │   W_0    │  F/4
    │ X_0 │ X_1 │ X_2 │ X_3 │    ├──────────┤
    └─────┴─────┴─────┴─────┘    │   W_1    │  F/4
       F/4 each                   ├──────────┤
                                  │   W_2    │  F/4
                                  ├──────────┤
                                  │   W_3    │  F/4
                                  └──────────┘
                                       D

  每卡本地算:
    GPU 0: Y_0 = X_0 · W_0  ∈ [B*S, D]  (是"部分和",不是最终结果)
    GPU 1: Y_1 = X_1 · W_1  ∈ [B*S, D]
    GPU 2: Y_2 = X_2 · W_2  ∈ [B*S, D]
    GPU 3: Y_3 = X_3 · W_3  ∈ [B*S, D]

  最终结果:
    Y = Y_0 + Y_1 + Y_2 + Y_3  ← 需要 AllReduce！

核心：输入是分片的（正好接 Column Parallel 的输出）、输出是部分和——必须 AllReduce 才能得到完整结果。

4.3 Column + Row 配对：一层 MLP 只需要一次 AllReduce

        [B, S, D]                       [B, S, D]
     "复制" 输入                       "复制" 输出
        每卡都有                          每卡都有
           │                                 ▲
           ▼                                 │
    ┌─────────────┐                   ┌──────┴──────┐
    │  W_up       │                   │  W_down     │
    │ Column      │                   │ Row Parallel│
    │ Parallel    │                   │             │
    │ [D → F]     │                   │ [F → D]     │
    └──────┬──────┘                   └──────▲──────┘
           │                                 │
           ▼                                 │
    [B, S, F/4]                       [B, S, D]  (部分和)
    "分片" 状态                       每卡有一个,需 AllReduce
    每卡 F/4                                 ▲
           │                                 │
           │     ┌────────────┐             │
           └────►│ activation │─────────────┘
                 │   (SiLU)   │
                 └────────────┘
                 每卡独立算,无通信

  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
  整个 MLP 只有一次 AllReduce (发生在 W_down 之后)
  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

为什么 Column + Row 是"免费"的：中间那个 [B, S, F/4] 的分片状态天然衔接——从 Column 出来就是分片,直接送进 Row 不需要额外通信。如果反过来 Row + Column,中间就需要一次 AllGather + 一次 AllReduce,通信翻倍。

Attention 也是同一套配对：

                 Attention TP 分卡示意（N=4）

  输入 X ∈ [B, S, D]  (复制)
                │
                ▼
        ┌───────────────┐
        │  W_qkv        │  Column Parallel
        │  [D → 3*D]    │  按 head 切
        │               │
        │ 切成 4 份:    │
        │  head 0-15  │ head 16-31 │ head 32-47 │ head 48-63
        │  ──────────┼─────────────┼────────────┼────────────
        │  GPU 0     │  GPU 1      │  GPU 2     │ GPU 3
        └───────┬───────┘
                ▼
      每卡持有 (Q, K, V) ∈ [B, S, D/4]
                │
                ▼
        ┌───────────────┐
        │  Attention    │
        │  本地计算!    │  每卡 16 个 heads 独立算
        │  无通信       │
        └───────┬───────┘
                ▼
      每卡持有 O_local ∈ [B, S, D/4]
                │
                ▼
        ┌───────────────┐
        │  W_o          │  Row Parallel
        │  [D → D]      │  按 head 切输入
        │               │
        │  Y = O · W_o  │  每卡本地算,得到 [B, S, D] (部分和)
        └───────┬───────┘
                │
                ▼ AllReduce
      Y = sum(Y_local)  ∈ [B, S, D]  (完整,复制状态)

一层 Transformer 总共 2 次 AllReduce（MLP 一次 + Attention 一次）。Forward 2 次,Backward 又 2 次,一步 4 次。这就是前面通信量表的来源。

4.4 为什么 TP 必须机内

用 Llama 70B 算一笔账,一次 AllReduce 传的是 [B, S, D] 大小的激活:

单次 AllReduce 数据量 = B × S × D × 2 bytes (bf16)
                    = 1 × 8192 × 8192 × 2
                    = 128 MB

每层 4 次 (forward 2 + backward 2)
80 层  = 80 × 4 × 128 MB = 40 GB/step

延迟对比:
  NVLink Gen4 (H100): 900 GB/s  →  40 GB / 900 = 44 ms
  InfiniBand HDR:      50 GB/s  →  40 GB / 50  = 800 ms

  差 18 倍!  跨机 TP 会让通信占 50%+ 时间

结论：TP 只能机内。这也是为什么一台 8×H100 机器就是 TP=8 的天然上限。

五、Sequence Parallel：把 LayerNorm 也切了

TP 存在一个浪费：**元素级操作（LayerNorm、Dropout、Residual）**在 TP 下每卡都在算完整 [B, S, D]——没有并行加速，纯浪费显存和算力。

5.1 SP 的核心 insight

element-wise op 沿 sequence 维度切开是独立的（不像 attention 需要跨 sequence 交互）。所以可以按 S 切。

           TP-only（浪费）:                   TP + SP（省显存）:

  ┌─────────────────────────┐           ┌─────────────────────────┐
  │ 前一个 Attention 块      │           │ 前一个 Attention 块      │
  │ 输出: [B, S, D] (复制)   │           │ 输出: [B, S/4, D] (分片) │
  └──────────┬──────────────┘           └──────────┬──────────────┘
             │                                     │
             ▼ (无通信,已复制)                     ▼ AllGather(sp)
                                                   拉齐 sequence
  ┌─────────────────────────┐           ┌─────────────────────────┐
  │ LayerNorm               │           │ LayerNorm               │
  │ 每卡算完整 [B, S, D]    │           │ 每卡算 [B, S/4, D]      │
  │ 4 张卡在重复干活!       │  ◄──────  │ 4 张卡真正并行           │
  └──────────┬──────────────┘           │ (但先要 AllGather!)      │
             │                          └──────────┬──────────────┘
             ▼                                     ▼
      进入 Attention TP 区                  进入 Attention TP 区

  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
  之后 Attention 输出 → AllReduce → LayerNorm → ...
  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

看起来 SP 多了一次 AllGather,通信不是更多了吗?

5.2 关键魔法：把 AllReduce 拆成两半

用第二节的恒等式:

AllReduce = ReduceScatter + AllGather

原来的 TP:

Attention 输出 (部分和) ──[AllReduce]──► LayerNorm 输入 (复制,完整 S)
                            ↑
                     一次 collective

改成 TP+SP:

Attention 输出 (部分和) ──[ReduceScatter over S]──► LayerNorm 输入 (完整 sum, S/4 分片)
                                                              │
                                                              ▼
                                                     LayerNorm (本地算 [B, S/4, D])
                                                              │
                                                              ▼
                                             ──[AllGather over S]──►  下一个 TP 区输入
                                                              (完整 S, 复制)

通信量完全没变! AllReduce 拆成 ReduceScatter + AllGather 两半,中间插入了本地 LayerNorm——LayerNorm 的激活值只需要存 [B, S/4, D] 而不是 [B, S, D]。

5.3 SP 的收益

一层 Transformer 的激活值显存（在 element-wise op 处）:

  纯 TP:  4 × B × S × D × 2 bytes   (LayerNorm × 2 + Dropout × 2)
  TP+SP:  4 × B × S/N × D × 2 bytes

N=4 时: 显存下降到 1/4 (仅 element-wise 部分)

整层激活值构成(不 checkpointing):
  ┌──────────────────────────────────────────┐
  │ Attention: QKV 中间, softmax matrix       │  ← FlashAttention 已优化
  │ ────────────────────────────────────────  │
  │ FFN: [B,S,F] 中间态 (超大!)                │  ← TP 已解决
  │ ────────────────────────────────────────  │
  │ element-wise: LayerNorm, Dropout, Residual│  ← SP 解决这块!
  └──────────────────────────────────────────┘

面试锚点：SP 通信量和 TP 一样,但显存节省显著。Megatron-LM 现代实现里 TP 和 SP 几乎总是一起用。

六、Context Parallel：把 Attention 的序列也切了

长上下文（32K+）时,Attention 的 [B, S, S] 矩阵会爆显存。CP 把 sequence 维度切开,让 attention 也分布式做。

6.1 为什么 CP 不能简单切 S

Attention 有个诅咒:每个 token 需要 attend 所有其他 token。如果每卡只有一段 sequence,单卡本地算 Q·K^T 只能拿到 attention matrix 的一小块。

序列切 4 段（CP=4）,每卡只有 S/4 的 Q,K,V:

  完整 Attention Matrix (S × S):
        Q_0 ── K_0
        Q_1 ── K_1
        Q_2 ── K_2      ← 我们需要每个 Q_i 和所有 K_j 交互
        Q_3 ── K_3      ← 但每卡只有一段 K!

  ┌───────────────────────────────┐
  │ Q_0 K_0 │ Q_0 K_1 │ Q_0 K_2 │ Q_0 K_3│
  │─────────┼─────────┼─────────┼────────│
  │ Q_1 K_0 │ Q_1 K_1 │ Q_1 K_2 │ Q_1 K_3│
  │─────────┼─────────┼─────────┼────────│
  │ Q_2 K_0 │ Q_2 K_1 │ Q_2 K_2 │ Q_2 K_3│
  │─────────┼─────────┼─────────┼────────│
  │ Q_3 K_0 │ Q_3 K_1 │ Q_3 K_2 │ Q_3 K_3│
  └───────────────────────────────┘
      ▲
      本地能算的只有对角块！其他块的 K/V 都在别的卡上

6.2 Ring Attention: 让 K,V 循环转一圈

核心想法:Q 留在原地,让 K,V 沿环形拓扑传递,每转一圈每张卡都见过所有 K,V。

Ring Attention 时序图（4 张卡）:

  时间 →

  卡 0:  ┌──────┬──────┬──────┬──────┐
         │ Q₀·K₀│ Q₀·K₃│ Q₀·K₂│ Q₀·K₁│  ← Q_0 本地不动,轮流用 K_0,K_3,K_2,K_1
         └──────┴──────┴──────┴──────┘
         同时: K,V 从上一张卡收进来,发给下一张卡

  卡 1:  ┌──────┬──────┬──────┬──────┐
         │ Q₁·K₁│ Q₁·K₀│ Q₁·K₃│ Q₁·K₂│
         └──────┴──────┴──────┴──────┘

  卡 2:  ┌──────┬──────┬──────┬──────┐
         │ Q₂·K₂│ Q₂·K₁│ Q₂·K₀│ Q₂·K₃│
         └──────┴──────┴──────┴──────┘

  卡 3:  ┌──────┬──────┬──────┬──────┐
         │ Q₃·K₃│ Q₃·K₂│ Q₃·K₁│ Q₃·K₀│
         └──────┴──────┴──────┴──────┘

  每一步:
    - 用当前手上的 K,V 和本地 Q 算 attention 一块
    - 同时把当前 K,V 传给下一张卡（异步）
    - N 步后每张卡都用完了所有的 K,V

  关键: 通信和计算完全 overlap（send/recv 是异步的）
        只要一步计算时间 >= 一次 K,V 传输时间,通信就被完全隐藏

6.3 online softmax 让分块 attention 数学正确

朴素的分块会遇到 softmax 的问题——softmax 的分母需要看到所有值才能算。FlashAttention 用的 online softmax 解决了这个:每次算一块就更新 running max 和 running sum,最后归一化。

online softmax 状态维护（每张卡本地）:

  m: running max (每个 Q 位置一个)
  l: running sum-of-exp (每个 Q 位置一个)
  O: running output (每个 Q 位置一个 D 维向量)

  每收到一块 K,V:
    1. 算 S_new = Q · K^T
    2. 更新 m_new = max(m_old, max(S_new))
    3. 用新旧 max 差重新 rescale l 和 O
    4. 累加 exp(S_new - m_new) 到 l
    5. 累加 exp(S_new - m_new) · V 到 O

  N 步后:
    最终 O / l 就是正确的 attention 输出

Ring Attention = Ring 传 K,V + online softmax + 通信/计算 overlap——这就是 Gemini 1M context 背后的技术。

七、FSDP / ZeRO：ailzhang 留的坑填一下

FSDP 是最重要的"训练侧显存优化"。ailzhang 没画的这一块,我详细讲。

7.1 ZeRO 三级递进

核心 insight：DP 下每张卡都持有完整参数+梯度+优化器状态是巨大浪费——DP 组内的 N 张卡都在存同样的东西,只是各自算不同的 batch。

纯 DP 每卡的显存构成（Adam 优化器,bf16 训练）:

  ┌─────────────────────────────────────────────┐ 每卡都有
  │ 参数 (bf16)     : params × 2 bytes           │ 完整一份
  │ 梯度 (bf16)     : params × 2 bytes           │ ← 重复!
  │ Adam m (fp32)   : params × 4 bytes           │
  │ Adam v (fp32)   : params × 4 bytes           │
  │ master (fp32)   : params × 4 bytes           │
  │ ────────────────────────────────────────────│
  │ 合计            : 16 × params                │
  └─────────────────────────────────────────────┘

  70B 模型 = 70B × 16 = 1.1 TB/卡  ← OOM!

ZeRO 分三级切开这些"重复存"的东西:

ZeRO-1: 切优化器状态

  参数     : 每卡都有 (复制)         2 GB (对应 1B params)
  梯度     : 每卡都有 (复制)         2 GB
  优化器   : 按 N 切! 每卡 1/N        12/N GB  ← 省了！

  N=4: 每卡 (2+2+3) = 7 GB (vs 原来 16 GB)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

ZeRO-2: 优化器 + 梯度

  参数     : 每卡都有 (复制)         2 GB
  梯度     : 按 N 切! 每卡 1/N        2/N GB
  优化器   : 按 N 切! 每卡 1/N        12/N GB

  N=4: 每卡 (2 + 0.5 + 3) = 5.5 GB

  实现: AllReduce grad → 改成 ReduceScatter grad
        (通信量不变,但每卡只存 1/N 梯度)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

ZeRO-3 / FSDP: 三样都切

  参数     : 按 N 切! 每卡 1/N        2/N GB
  梯度     : 按 N 切! 每卡 1/N        2/N GB
  优化器   : 按 N 切! 每卡 1/N        12/N GB

  N=4: 每卡 (0.5 + 0.5 + 3) = 4 GB

  代价: forward/backward 时需要临时 AllGather 参数
        通信量 ~1.5x DP

7.2 FSDP 的运行时时间线

FSDP 最难理解的是运行时——每次前向/反向传播都要 AllGather + discard 参数。看时间线:

                    FSDP forward pass 时间线（3 层示意）

  时间 →

  layer 0:                    ┌─ AllGather 参数 ─┐
                              │   ┌── compute ──┐│
                              │   │   forward   ││
                              │   └─────────────┘│
                              │   ┌ discard 参数┘
                              │
  layer 1:                    │  ┌─ AllGather 参数 ─┐  ← 与 layer0 compute overlap!
                              │  │   ┌── compute ──┐│
                              │  │   │   forward   ││
                              │  │   └─────────────┘│
                              │  │   ┌ discard 参数┘

  layer 2:                        ┌─ AllGather 参数 ─┐
                                  │   ┌── compute ──┐│
                                  │   │   forward   ││
                                  │   └─────────────┘│

  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
  overlap 是灵魂: 当前 layer 计算时,下一层的 AllGather 并行进行
  只要 compute >= 通信,通信被完全隐藏
  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

                    FSDP backward pass 时间线

  layer 2 backward:  ┌─AllGather─┐ ┌─backward─┐ ┌─ReduceScatter grad─┐
                                      │                │
  layer 1 backward:                   ├─AllGather─┐    │  ┌─backward─┐ ...
                                                       │  │
                                                       └────► 拉下一层参数
  layer 0 backward:                                          ...

  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
  backward 每层做三件事:
    1. AllGather 参数 (临时拉完整,和 forward 时形状一样)
    2. 计算 backward
    3. ReduceScatter 梯度 (分散到各卡)
  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

通信量分析:

一次完整 iteration (forward + backward + step):

  FSDP:
    forward:  L × AllGather(params/N)  ≈ 一次 params
    backward: L × AllGather(params/N)  ≈ 一次 params
              L × ReduceScatter(grad/N) ≈ 一次 params
    ─────────────────────────────────
    合计: 3 × params

  纯 DP:
    backward: AllReduce(gradients) = 2 × params (Ring AllReduce 系数)
    ─────────────────────────────────
    合计: 2 × params

  FSDP 通信量 = 1.5x DP,但显存降到 params/N

面试锚点：ailzhang 说 FSDP 不影响 local tensor shape——从 "compute 时看到什么" 的视角是对的（AllGather 后确实看到完整参数）。但从 "存什么" 的视角,FSDP 每张卡只存 1/N 参数。两个视角都对,不要混淆。

7.3 FSDP × TP 怎么搭：HSDP

现代大模型训练常用 HSDP (Hybrid Sharded Data Parallel)：

                      HSDP 拓扑示意

  ┌────────────────────┐    ┌────────────────────┐
  │   机器 A            │    │   机器 B            │
  │  ┌───┬───┬───┬───┐ │    │  ┌───┬───┬───┬───┐ │
  │  │G0 │G1 │G2 │G3 │ │    │  │G0 │G1 │G2 │G3 │ │
  │  ├───┼───┼───┼───┤ │    │  ├───┼───┼───┼───┤ │
  │  │G4 │G5 │G6 │G7 │ │    │  │G4 │G5 │G6 │G7 │ │
  │  └───┴───┴───┴───┘ │    │  └───┴───┴───┴───┘ │
  │  NVLink (机内)     │    │  NVLink (机内)     │
  │  ┌──────────────┐  │    │  ┌──────────────┐  │
  │  │ TP=8         │  │    │  │ TP=8         │  │
  │  │ 每层参数在   │  │    │  │ 每层参数在   │  │
  │  │ 8 卡间切分   │  │    │  │ 8 卡间切分   │  │
  │  └──────────────┘  │    │  └──────────────┘  │
  └──────────┬─────────┘    └─────────┬──────────┘
             │                        │
             └────── IB (机间) ───────┘
                    FSDP 跨机 shard
                    参数在机器间切分
                    (每机器只存 1/M 参数副本)

  两层 shard 结构:
    - 机内: TP=8 (Megatron 方式,按 head/F 切权重)
    - 机间: FSDP=M (按 DP 方式,跨机 shard 参数)

  总参数副本数 = 1 份 (完全分散)
  单卡负担 = params / (TP × FSDP)

取舍：

FSDP-only（组内组间都 shard）：显存最省,但跨机通信压力大
HSDP（组内 TP,组间 FSDP）：现代 H100 集群主流,机内 TP 高带宽 + 机间 FSDP 分显存
ZeRO-2 只切梯度/优化器：显存节省不如 FSDP,但通信压力小

八、Pipeline Parallel 和它的 Bubble

PP 的核心 insight：把 L 层切成 K 段,每段放一台机器,forward 时数据流过 K 台机器。

8.1 朴素 PP 的 bubble 问题

                朴素 PP 时间线（4 段,一次一个 batch）

  时间 →  ═════════════════════════════════════════

  GPU 0: ┃F0┃                              ┃B0┃
  GPU 1:     ┃F1┃                      ┃B1┃
  GPU 2:         ┃F2┃              ┃B2┃
  GPU 3:             ┃F3┃      ┃B3┃

           ├──stage 内 forward──┤├─backward-┤
           GPU 大部分时间在等!

  bubble 比例 = (K-1)/K = 3/4 = 75% 空闲！

8.2 GPipe → 1F1B → Interleaved 1F1B

GPipe（Google, 2018）：把 batch 切成 M 个 micro-batch,像流水线一样喂：

                GPipe 时间线（K=4, M=8）

  时间 →

  GPU 0: F0 F1 F2 F3 F4 F5 F6 F7                  B7 B6 B5 B4 B3 B2 B1 B0
  GPU 1:    F0 F1 F2 F3 F4 F5 F6 F7             B7 B6 B5 B4 B3 B2 B1 B0
  GPU 2:       F0 F1 F2 F3 F4 F5 F6 F7       B7 B6 B5 B4 B3 B2 B1 B0
  GPU 3:          F0 F1 F2 F3 F4 F5 F6 F7 B7 B6 B5 B4 B3 B2 B1 B0

  bubble 比例 = (K-1)/(K+M-1) = 3/11 ≈ 27%
  M=32 时  = 3/35 ≈ 9%
  M=64 时  = 3/67 ≈ 5%

  问题: 所有 forward 完再做 backward → 激活值要存 M 次！
        M=64 时激活值显存爆炸

1F1B (One Forward One Backward, PipeDream)：交替做 forward 和 backward,激活值只存 K 次:

                1F1B 时间线（K=4, M=8）

  时间 →

  GPU 0: F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7
  GPU 1:    F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7
  GPU 2:       F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7
  GPU 3:          F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7

  bubble 比例 = 和 GPipe 相同,但激活值显存降到 O(K) 不是 O(M)!

  1F1B 稳态: 每张卡在 F(i) 和 B(i-K+1) 之间交替
             一旦 backward 开始,forward 一步就 backward 一步

Interleaved 1F1B：每张卡持有多个不连续的 layer chunk（如 GPU 0 持有 L0-L10 和 L40-L50）。让流水线更细,bubble 进一步降低。代价是通信次数翻倍。

8.3 DeepSeek V3 的 DualPipe

DualPipe 同时跑两条 pipeline,一条正向一条反向,让它们的通信互相 overlap：

                DualPipe 概念示意（简化）

  时间 →

  GPU 0: F(正向 batch A) │ B(反向 batch B) │ F(正向 A) │ B(反向 B)
                            ↑                           ↑
                       两条 pipeline 的通信在此处 overlap

  GPU 1: 类似
  GPU 2: 类似
  GPU 3: 类似

  精妙之处:
    - 一条 pipeline 的 backward 通信 (send activations)
      与另一条的 forward compute overlap
    - MoE 的 AllToAll 通信也塞进 pipeline 空闲气泡里
    - bubble 几乎归零

  代价: 显存翻倍 (两条 pipeline 同时存激活)

面试锚点：Pipeline bubble 无法完全消除的根源是依赖关系——第 k 层的 forward 必须等第 k-1 层输出,backward 必须等第 k+1 层梯度传回。1F1B/DualPipe 都是通过精巧调度让 GPU 不空闲,但第一个 micro-batch 传到最后一层前,前面几张卡必然空闲（warm-up bubble）。

九、MoE Expert Parallel 深度展开

MoE 的 EP 是当代最复杂的并行。用 8 个 expert × 4 张卡（EP=4）为例:

9.1 Expert 在卡上的分布

                Expert Parallel 布局（8 experts, EP=4）

  GPU 0: E0, E4          ← 每卡持有 2 个 expert
  GPU 1: E1, E5
  GPU 2: E2, E6
  GPU 3: E3, E7

  每个 expert 是一个独立的 FFN 权重（W_up + W_down + W_gate）
  Llama-3-70B 那种 dense 模型没有,DeepSeek/Mixtral/Qwen3 才有

9.2 MoE forward 完整流程

                MoE Forward 6 步走

  Step 1: Router 给每个 token 打分（本地计算,无通信）
  ─────────────────────────────────────────────────────
    GPU 0: 拿到 tokens [t0..t99]
    经过 router MLP → 每个 token 的 8 个 expert 的分数
    Top-2 路由: t0 → E1,E3; t1 → E0,E5; ...


  Step 2: Local Permutation（本地重排 token）
  ─────────────────────────────────────────────────────
    按目标 expert 分组:
    GPU 0 本地重排后:
      Bucket_E0: [t1, t5, t9, ...]        ← 目标 E0 (在 GPU 0 上,本地)
      Bucket_E1: [t0, t3, t7, ...]        ← 目标 E1 (在 GPU 1 上,要发)
      Bucket_E2: [t2, t8, ...]
      Bucket_E3: [t0, t11, ...]
      Bucket_E4: [...]                    ← 目标 E4 (在 GPU 0 上,本地)
      Bucket_E5: [...]
      Bucket_E6: [...]
      Bucket_E7: [...]


  Step 3: Dispatch AllToAll（第一次通信!）
  ─────────────────────────────────────────────────────
        ┌─────────────────────────────────────┐
    GPU 0│──[Bucket_E1]──►GPU 1                │
        │──[Bucket_E2]──►GPU 2                │
        │──[Bucket_E3]──►GPU 3                │
        │──[Bucket_E0,E4]─保留本地             │
        │──[Bucket_E5]──►GPU 1                │
        │──[Bucket_E6]──►GPU 2                │
        │──[Bucket_E7]──►GPU 3                │
        │                                     │
    ... 每张卡同理,N² 条链路同时打开          │
        └─────────────────────────────────────┘

    通信后每张卡收到该卡 expert 该处理的所有 token:
      GPU 0 (跑 E0, E4): [来自所有卡的 E0 tokens] + [来自所有卡的 E4 tokens]
      GPU 1 (跑 E1, E5): 类似


  Step 4: Expert 计算（本地,无通信）
  ─────────────────────────────────────────────────────
    每张卡的每个 expert 独立算 FFN
    GPU 0:
      E0_output = FFN_E0(E0_tokens)
      E4_output = FFN_E4(E4_tokens)

    ⚠️ 每卡收到的 token 数不同 → 负载不均衡,拖尾严重


  Step 5: Combine AllToAll（第二次通信!）
  ─────────────────────────────────────────────────────
    把 expert 输出送回原来的 token 位置
    和 Step 3 完全对称,但方向反了


  Step 6: Un-permutation + 加权组合
  ─────────────────────────────────────────────────────
    每个 token 收到 2 个 expert 的输出（top-2）
    按 router 打分加权组合:
      out = w_E1 · out_E1 + w_E3 · out_E3

一层 MoE 有 2 次 AllToAll。80 层就是 160 次——这就是为什么 MoE 训练 AllToAll 是绝对瓶颈。

9.3 EP × TP 组合：万亿参数模型的典型拓扑

DeepSeek V3 (671B, MoE) 的实际配置:

                  DeepSeek V3 训练拓扑（简化）

  ┌─────────────────────────────────────────────────┐
  │  Cluster: N 台机器,每台 8 卡                     │
  │                                                  │
  │  一层 Transformer 的 forward:                    │
  │                                                  │
  │  ┌──────────────┐                                │
  │  │  Attention   │  ← 用 TP 切 head              │
  │  │  TP=1 or 2   │    (Llama 用 8,DeepSeek 少)   │
  │  └──────┬───────┘                                │
  │         │                                        │
  │  ┌──────▼───────┐                                │
  │  │    Router    │  ← 本地计算                    │
  │  └──────┬───────┘                                │
  │         │                                        │
  │  ┌──────▼───────┐                                │
  │  │   MoE FFN    │  ← EP=64 (跨机!)              │
  │  │   256 experts│    每台机器持有一部分 experts │
  │  │              │    AllToAll 走 IB              │
  │  └──────────────┘                                │
  │                                                  │
  │  额外:                                           │
  │    - PP=16 (Pipeline 分层)                       │
  │    - DualPipe (让 AllToAll 塞进 PP 气泡)         │
  │    - FSDP 组内 shard 参数                        │
  └─────────────────────────────────────────────────┘

十、推理侧并行：TP/EP/PP 的差异

前面全在讲训练。推理侧我的主场,补关键差异。

10.1 推理和训练的核心差异

                训练 vs 推理 关键差异

  维度         │ 训练              │ 推理
  ─────────────┼───────────────────┼──────────────────
  主导阶段     │ forward + backward│ 只有 forward
  Batch 大小   │ GBS 512-2048 稳定 │ 动态 batch,变化剧烈
  关键指标     │ Throughput (MFU)  │ Latency (TTFT + TPOT)
  通信容忍度   │ 高（可 overlap）  │ 低（延迟敏感）

结论：推理下通信开销更难被隐藏（decode 阶段计算量小,compute 遮不住通信）。

10.2 推理下的 TP：为什么止步 TP=8

              训练 TP AllReduce (可 overlap):

  layer i forward compute ┌───────────────────┐
                          │                   │
                          └──────┬────────────┘
  AllReduce               ┌────▼─────┐
                          └──────┬───┘
  layer i+1 fwd compute        ┌─▼─────────────┐  ← 与 AllReduce overlap
                               └───────────────┘
                          总时间 = max(compute, comm)

  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
              推理 TP AllReduce (无法 overlap):

  decode step (只算 1 token 的 forward):
    layer i compute  ┌──┐            ← 极短(1 token 太少)
                      └┬─┘
    AllReduce         ┌▼──────────┐  ← 通信主导!
                      └───────────┘
    layer i+1 compute            ┌──┐
                                  └──┘

  总时间 = compute + comm  ← 通信不能藏!

Decode 阶段每步只算 1 token,compute 时间极短,AllReduce 通信占大头
推理 TP 通常最多 TP=8（一台机器内 NVLink）——TP=16 跨机会大幅掉 TPOT
MQA/GQA 帮了大忙——KV heads 少,TP 对 KV 的复制开销小

10.3 推理下的 EP：DP Attention + EP FFN

MoE 推理最有意思的是 Attention 走 DP、FFN 走 EP 的组合:

                  DP Attention + EP FFN 布局

  8 张卡,MoE 32 experts

  ┌──────────────────────────────────────────────────┐
  │  各卡完全独立处理不同请求（DP Attention）           │
  │                                                    │
  │  GPU 0: [请求 A] → Attention 本地  → 拿到 hidden_A│
  │  GPU 1: [请求 B] → Attention 本地  → 拿到 hidden_B│
  │  GPU 2: [请求 C] → Attention 本地  → 拿到 hidden_C│
  │  GPU 3: [请求 D] → Attention 本地  → 拿到 hidden_D│
  │  ...    (每卡独立,无跨卡通信!)                   │
  │                                                    │
  │  ↓ 到 MoE FFN 时                                   │
  │                                                    │
  │  8 张卡共享 32 experts (每卡 4 个 expert)         │
  │  用 AllToAll 路由 hidden state                     │
  │                                                    │
  │  GPU 0: 持有 E0-E3                                 │
  │  GPU 1: 持有 E4-E7                                 │
  │  ...                                               │
  │                                                    │
  │  AllToAll → 每张卡收到给自己 expert 的 tokens     │
  │             (来自所有请求)                          │
  │  本地 FFN → AllToAll 送回原位置                    │
  └──────────────────────────────────────────────────┘

  为什么这个组合好:
    - Attention 完全无跨卡通信 (每请求独立)
    - 只有 FFN 层需要 AllToAll,通信频次减半
    - 每个请求的 KV Cache 完全在自己卡上,不用分享
    - 天然支持 Prefix Cache (KV 就地存,不跨卡)

10.4 推理独有的并行：DP 副本 + Prefix Cache

推理有个训练没有的维度——多副本 DP 服务:

                  推理集群拓扑

                       ┌──────────┐
                       │  Router  │  按 prompt hash 路由
                       └────┬─────┘
                            │
             ┌──────────────┼──────────────┐
             │              │              │
        ┌────▼───┐    ┌─────▼──┐    ┌────▼────┐
        │副本 1  │    │副本 2  │    │副本 3   │
        │        │    │        │    │         │
        │TP=8    │    │TP=8    │    │TP=8     │
        │Prefix  │    │Prefix  │    │Prefix   │
        │Cache   │    │Cache   │    │Cache    │
        └────────┘    └────────┘    └─────────┘

  路由策略: 同一 prompt 前缀 → 同一副本
            利用副本的 Prefix Cache (KV Cache 复用)

  好处:
    - 无梯度同步 (推理没有梯度)
    - 副本之间完全独立
    - 前缀命中率高 (system prompt 共享率 >90%)

这是推理独有的并行拓扑——vLLM 三年演进史里详细讲过。

10.5 为什么推理很少用 SP/CP

SP 为省 LayerNorm 激活显存,推理没有 backward 不需要存激活
CP 为训练超长序列,推理长序列走 FlashAttention + PagedAttention 就够了

Pipeline Parallel 推理里极少用——bubble 在流式服务里消除不掉,只在"单机装不下的超大模型"时不得不 PP。

十一、选型决策树

给一个可以照抄的选型流程:

                    分布式选型决策树

  开始
    │
    ▼
  ┌─────────────────────────┐
  │ 单卡能装下模型吗？        │
  └───┬─────────────────────┘
      │
   能 │       ┌── DP + ZeRO-1 起步（省显存不掉性能）
      ▼───────┤
      │       └── 数据够多就上,不需要复杂并行
   不能
      ▼
  ┌─────────────────────────┐
  │ 有 NVLink 机内互联吗？    │
  └───┬─────────────────────┘
      │
   有 │       ┌── 机内 TP=8 起步（Column+Row 配对）
      ▼───────┤
      │       └── 可能配上 SP 省激活值显存
   没有
      ▼───────── 只能纯 PP（少数不幸的场景）


  ┌─────────────────────────┐
  │ 加上 TP=8 后单机能装下吗？│
  └───┬─────────────────────┘
      │
   能 │       ┌── 机间 DP 或 FSDP
      ▼───────┤
      │       └── 现代主流: HSDP (机内 TP + 机间 FSDP)
   不能
      ▼───────── 加 PP,跨机 stage


  ┌─────────────────────────┐
  │ 是 MoE 吗?               │
  └───┬─────────────────────┘
      │
   是 ▼───────── 加 EP,机内先做,溢出再机间
      │        （DualPipe 让 AllToAll overlap PP 空闲）
      │
   否 ▼


  ┌─────────────────────────┐
  │ 序列超长(>32K)?          │
  └───┬─────────────────────┘
      │
   是 ▼───────── 加 SP + CP,和 TP 配合


  ┌─────────────────────────┐
  │ 显存还是不够?             │
  └───┬─────────────────────┘
      │
   是 ▼───────── 上 FSDP 或 ZeRO-3（代价 1.5x 通信）

典型 Llama-3 70B 训练:

8 台 H100 = 64 卡
每台 TP=8（NVLink）
机间 FSDP=8（IB）
无 PP,无 EP

典型 DeepSeek V3 671B (MoE) 训练:

数千卡集群
TP=1（MoE expert 本身分卡）
EP=64（跨节点）
PP=16 + DualPipe
FSDP 组内 shard 参数

典型 kimi-k2 万亿参数推理（推测）:

单请求 TP=8 + EP=32（几十张卡服务一个请求）
多副本 DP 复制
Prefix Cache + Continuous Batching

十二、面试锚点

Q1: TP 为什么必须走机内 NVLink?

A: TP 每层 4 次 AllReduce（fwd 2 + bwd 2）,每次约 B×S×D 级别。Llama 70B、B=1、S=8192、D=8192,一次 AllReduce ≈ 128 MB,80 层就是 40 GB/step。跨机 IB（50 GB/s）跑要 800ms,机内 NVLink（900 GB/s）只要 44ms。差 18 倍。TP 只能机内。

Q2: FSDP 和 ZeRO-3 有什么区别?

A: 概念上是同一件事——都把参数、梯度、优化器状态全按 DP 组切开。工程实现不同:ZeRO-3 是 DeepSpeed 的实现,FSDP 是 PyTorch 官方实现,通信调度和 hook 机制略有差异。等价的,选实现看框架偏好。

Q3: AllReduce 为什么能拆成 ReduceScatter + AllGather?这有什么意义?

A: 数学上 AllReduce = "reduce 到每卡 + 广播"。ReduceScatter 是"reduce 完每卡只拿一片",AllGather 是"每卡有一片聚成完整",合起来等于 AllReduce。意义:中间那一步"每卡只有一片"可以直接用于分布式优化器更新——这就是 ZeRO-2/FSDP 省显存的原理。SP 也用同样的等价变换,把 AllReduce 拆开后中间插入 element-wise op。

Q4: MoE 的 AllToAll 为什么这么头疼?

A: 三个原因: (1) 通信量大——每卡把 token 分成 N 份发出去再收回来,总量 O(B×S×D); (2) 不能用 ring——N² 条链路同时活跃,fabric 压力极大; (3) 延迟主导——路由必须完成才能算 FFN,不像 AllReduce 可以 overlap 到反向传播。DeepSeek V3 用 DualPipe 把 AllToAll 塞进 PP 空闲气泡是目前最先进的方案。

Q5: Ring AllReduce 为什么带宽不受卡数影响?

A: Ring 把 tensor 切成 N 块,ReduceScatter + AllGather 两阶段各 N-1 步。每一步每张卡同时发一块给下一张、收一块从上一张——所有链路同时工作。总通信量 2(N-1)/N × M ≈ 2M,和 N 无关。瓶颈从"带宽"变成"延迟"(O(N) 步)。1000+ 卡时切换到 Tree 或 Halving-Doubling。

Q6: FSDP 通信量比 DP 大,为什么大家还用?

A: 显存节省 >> 通信代价。FSDP 通信 1.5x DP,但显存降到 1/N。70B+ 大模型不用 FSDP 根本装不下（16 卡训 70B,不 shard 单卡 30GB+ 直接 OOM）。FSDP 通信可以和 compute overlap,实际吞吐差距远小于 1.5x。关键要有 NVLink——通信慢了 overlap 藏不住,FSDP 会崩。

Q7: 3D 并行（TP × PP × DP）的拓扑怎么设计?

A: "通信量大的走机内"原则:(1) TP=8 in machine (NVLink); (2) PP=K across machines within a pod (IB,延迟不敏感); (3) DP replicate across pods (AllReduce 可 overlap)。举例:1024 卡训 70B,TP=8, PP=4, DP=32——一台机器一个 TP 组,4 台机器一个 PP 组,32 组做 DP。

Q8（进阶）: Pipeline bubble 为什么无法完全消除?

A: 依赖关系——第 k 层的 forward 必须等第 k-1 层输出,backward 必须等第 k+1 层梯度传回。1F1B 通过交错 F 和 B 让 GPU 不空闲,但第一个 micro-batch 传到最后一层前,前面几张卡必然空闲（warm-up bubble）。DualPipe 用双向流水几乎消除,但代价显存翻倍。理论下界 K/(K+M),M 越大越好但激活值显存受限。

Q9（进阶）: 推理和训练的并行策略最大差异是什么?

A: 三点: (1) 推理没有 backward,SP/CP 不需要; (2) 推理延迟敏感,TP 通信藏不住,一般不超过 TP=8; (3) 推理独有 DP 副本 + Prefix Cache 组合——多副本、按前缀 hash 分发、共享 KV Cache。Pipeline Parallel 训练常用推理少用,bubble 在流式推理里消除不掉。

Q10（进阶）: 给你 1024 张 H100 训万亿参数 MoE,你会怎么切?

A: 参考 DeepSeek V3: (1) EP=64（跨节点专家并行,expert 数量决定）; (2) PP=8-16（Pipeline 分层）; (3) TP=1（MoE 本身分卡了,TP 反而增加通信）; (4) 组内 FSDP shard 参数; (5) 跨组 DP replicate。关键优化: DualPipe 让 AllToAll 和 PP 计算 overlap; hierarchical AllToAll 优先机内路由。

一句话记住这一篇

分布式 Transformer 的所有设计,本质在平衡三个约束（显存/算力/通信）——把通信量大的并行（TP/SP/EP）放到机内 NVLink,通信量小的（DP/PP）走机间 IB。理解 AllReduce = ReduceScatter + AllGather 这个恒等式就理解了 ZeRO/FSDP/SP 一系列优化的核心。训练侧和推理侧的并行策略在拓扑决策上有本质差异,推理独有 DP 副本 + Prefix Cache 这一维度,训练独有 PP 的 bubble 问题。

配套阅读:ailzhang 的 Visualizing Parallelism in Transformer 是最好的入门可视化,JAX Scaling Book 的 Transformer Accounting 是最好的单卡张量形状参考。我这个博客还有几篇相关: MoE 推理完整机制、vLLM 三年演进史、现代 LLM 推理引擎剖析、FlashAttention 三代演进。

分布式 Transformer 深度解析：从张量形状到通信拓扑

分布式 Transformer 深度解析：从张量形状到通信拓扑

一、先建立坐标系：一次 forward 到底有多贵

1.1 显存账（单卡装不下）

1.2 三种约束的取舍

二、通信原语：底层到底怎么走

2.1 四个基本 collective

2.2 Ring AllReduce：为什么带宽不受卡数影响

2.3 Tree AllReduce vs Halving-Doubling

2.4 AllToAll：MoE 的通信之王

三、六种并行 × 通信量对照表

四、Tensor Parallel：Column + Row 的免费配对

4.1 Column Parallel（切输出维度）

4.2 Row Parallel（切输入维度）

4.3 Column + Row 配对：一层 MLP 只需要一次 AllReduce

4.4 为什么 TP 必须机内

五、Sequence Parallel：把 LayerNorm 也切了

5.1 SP 的核心 insight

5.2 关键魔法：把 AllReduce 拆成两半

5.3 SP 的收益

六、Context Parallel：把 Attention 的序列也切了

6.1 为什么 CP 不能简单切 S

6.2 Ring Attention: 让 K,V 循环转一圈

6.3 online softmax 让分块 attention 数学正确

七、FSDP / ZeRO：ailzhang 留的坑填一下

7.1 ZeRO 三级递进

7.2 FSDP 的运行时时间线

7.3 FSDP × TP 怎么搭：HSDP

八、Pipeline Parallel 和它的 Bubble

8.1 朴素 PP 的 bubble 问题

8.2 GPipe → 1F1B → Interleaved 1F1B

8.3 DeepSeek V3 的 DualPipe

九、MoE Expert Parallel 深度展开

9.1 Expert 在卡上的分布

9.2 MoE forward 完整流程

9.3 EP × TP 组合：万亿参数模型的典型拓扑

十、推理侧并行：TP/EP/PP 的差异

10.1 推理和训练的核心差异

10.2 推理下的 TP：为什么止步 TP=8

10.3 推理下的 EP：DP Attention + EP FFN

10.4 推理独有的并行：DP 副本 + Prefix Cache

10.5 为什么推理很少用 SP/CP

十一、选型决策树

十二、面试锚点

一句话记住这一篇

评论