Sharded Matmul：分片矩阵乘的通用记号

前情说明：技术深度 D2。《读 JAX Scaling Book》系列第 2 篇。

上一篇 Roofline 分析回答了"一个 kernel 该多快"。这一篇回答一个更难的问题——一个被切开放在 N 张卡上的矩阵乘,到底该怎么算?

JAX Scaling Book 的 Ch3 给出了一套极其优雅的统一记号系统 [IX, JY]——不管你是 Megatron TP、FSDP、Context Parallel、Expert Parallel,还是任何没见过的新切法,这套记号都能描述。更妙的是,它把"分布式矩阵乘到底要插哪个 collective"压缩成了 4 种 case + 4 条规则。

这篇我在原书基础上做中文重构 + 加入 Megatron / FSDP / vLLM 推理的实际例子印证——让你学完能用它反过来看任何现有的并行方案。

一、为什么需要一套"统一记号"

1.1 现有讨论的混乱

你在网上看分布式训练文章,一定遇到过这些名字:

   Megatron-LM 系列:
     - Column Parallel / Row Parallel
     - Sequence Parallel
     - Tensor Parallel
     - Expert Parallel

   DeepSpeed / FSDP 系列:
     - ZeRO-1 / ZeRO-2 / ZeRO-3
     - FSDP full_shard / shard_grad_op / no_shard
     - HSDP

   JAX / PaxML 系列:
     - sharding_mesh
     - PartitionSpec
     - "replicated" / "partitioned"

   Gemini / DeepSeek 工程 blog:
     - DP / TP / PP / EP / SP / CP
     - 2D / 3D parallelism
     - DualPipe

这些概念互相重叠、互相矛盾。Column Parallel 到底和 Row Parallel 什么关系?FSDP 和 ZeRO-3 有什么区别?Sequence Parallel 为什么能和 TP"自然搭配"?

答案:它们本质上都是"把某个 tensor 的某个维度切到某个 mesh axis 上"——只是切的维度不同、起的名字不同。只要有一套记号能统一描述"谁切了谁",所有并行方案都变成这套记号的实例。

1.2 JAX Book 的解法:`[IX, JY]`

Google 的 JAX / XLA 团队用了一套"命名轴 + 下标标注"的记号。核心概念两个:

   概念 1: Device Mesh (设备网格)
   ─────────────────────────────────
     把 N 张卡抽象成一个多维网格,每一维有名字

     例: Mesh({'X': 4, 'Y': 2}) 表示 4x2 的二维网格
         ┌──────┬──────┬──────┬──────┐
         │ G0   │ G1   │ G2   │ G3   │  ← X 轴 (4 卡)
         ├──────┼──────┼──────┼──────┤
         │ G4   │ G5   │ G6   │ G7   │
         └──────┴──────┴──────┴──────┘
         ↑ Y 轴 (2 卡)


   概念 2: Sharding (分片规则)
   ─────────────────────────────────
     用下标标注每个逻辑维度切到了哪个 mesh axis 上

     A[I, J]        ← 复制 (每卡都有完整数据)
     A[IX, J]       ← I 维度切到 X 轴上
     A[IX, JY]      ← I 切 X,J 切 Y
     A[IXY, J]      ← I 维度切到 X 和 Y 两个轴 (flatten)
     A[IX, JX]      ← 禁止!一个 mesh axis 不能用两次

读法超简单:下标就是"这一维被哪些 mesh axis 切了"。没有下标就是复制。

1.3 一个小例子感受一下

   设定: Mesh({'X': 2, 'Y': 2}),4 卡组成 2x2 网格
         逻辑数组 A ∈ bf16[1024, 4096]

   情况 1: A[I, J]   (完全复制)
   ────────────────────────────────
     每卡持有完整 [1024, 4096]
     总存储 = 4 × 1024 × 4096 × 2 = 32 MB (浪费!)

   情况 2: A[IX, J]  (I 切到 X 轴)
   ────────────────────────────────
     每卡持有 [512, 4096]  (I/2=512)
     Y 轴上复制 (Y 组内 2 卡存一样的东西)
     总存储 = 4 × 512 × 4096 × 2 = 16 MB

   情况 3: A[IX, JY]  (I 切 X,J 切 Y)
   ────────────────────────────────
     每卡持有 [512, 2048]
     完全没有重复
     总存储 = 4 × 512 × 2048 × 2 = 8 MB (单份)

   情况 4: A[IXY, J]  (I 切到 X 和 Y 一起)
   ────────────────────────────────
     把 X 和 Y 看成一个 flatten 的 4 卡维度
     每卡持有 [256, 4096]  (I/4=256)
     总存储 = 4 × 256 × 4096 × 2 = 8 MB

三条规则记住:

下标越多 → 这个维度切得越细
不同维度可以用不同 mesh axis
同一个 mesh axis 不能切两个维度（A[IX, JX] 禁止）

二、分片矩阵乘的 4 种情况

现在到了最关键的部分。分布式矩阵乘的所有复杂度,都可以压缩成 4 种 case。

先理解一个铺垫概念:contracting dimension（收缩维度）。

   矩阵乘 A · B = C

     A[I, J]  ·  B[J, K]  →  C[I, K]
         ↑          ↑
       I 是       J 是
     non-contract  contract
      (保留)      (收缩, 内积求和)

   直觉:
     - non-contract 维度 → 决定输出形状,各卡独立
     - contract 维度 → 需要求和,这里是通信发生的地方!

所以关键问题是:A 和 B 的 contract 维度 J 有没有被切?如果切了,切在哪个 mesh axis 上?

Case 1: 都没切 contract 维度 → 零通信

   A[IX, J]  ·  B[J, KY]  →  C[IX, KY]
         ↑         ↑
       I 切 X     K 切 Y
         J 都是复制的 ✓

   直觉:
     每张卡本地都有完整的 J (因为没切)
     本地算完,结果天然按 X 和 Y 切好
     ──────────────────
     零通信! 这是最理想情况

面试锚点:为什么 FSDP 用 AllGather 把参数拉完整?就是为了让 contract 维度不切——变成 Case 1,零通信地做 matmul。

Case 2: 一边切了 contract 维度 → AllGather

   A[I, JX]  ·  B[J, K]  →  C[I, K]
         ↑        ↑
       J 切了 X   J 没切 (或两边 J 切在不同 axis)

   问题:
     A 每张卡只有 J 的一个 slice
     但 matmul 需要对完整 J 求和!
     ──────────────────
     解法: 先 AllGather A → 把 J 维度拉完整

   具体操作:
     Step 1: AllGather_X A[I, JX] → A[I, J]
     Step 2: A[I, J] · B[J, K] → C[I, K]

通信代价（单向环形 AllGather）:

   T_AllGather = V / W_ici
                 ↑       ↑
            数组字节数   带宽

   注意: 不依赖切几张卡! (bandwidth-bound regime)

Case 3: 两边都切 contract 维度在同一 axis → AllReduce

   A[I, JX]  ·  B[JX, K]  →  C[I, K]
         ↑        ↑
       J 都切了,且切在同一个 axis X 上

   直觉:
     每张卡能本地算一个"部分和"
     但只是全局求和的一部分
     ──────────────────
     解法: 本地 matmul 得到 partial sum,再 AllReduce 求和

   具体操作:
     Step 1: A[I, JX] ·_LOCAL B[JX, K] → C[I, K] {UX}
                                              ↑
                                        unreduced along X
     Step 2: AllReduce_X C[I, K]{UX} → C[I, K]

{UX} 这个"unreduced"记号是 JAX Book 的小创新——明确标出"这是一个还没求和完的中间状态"。

通信代价:

   T_AllReduce = 2 × V / W_ici    (是 AllGather 的 2 倍!)

为什么是 2 倍?因为 AllReduce = ReduceScatter + AllGather（这个恒等式在分布式 Transformer 深度解析里讲过）。

Case 4: 两边 non-contract 维度切在同一 axis → 先 AllGather

   A[IX, J]  ·  B[J, KX]  →  C[IX, KX]  ← 禁止!
         ↑        ↑          ↑
         X        X         X 出现两次,非法

   问题:
     结果 C 需要 I 和 K 两个维度都按 X 切
     但一个 mesh axis 只能切一个维度
     ──────────────────
     解法: 先把其中一个 AllGather 掉

   选项 A:
     AllGather_X A[IX, J] → A[I, J]
     A[I, J] · B[J, KX] → C[I, KX]

   选项 B:
     AllGather_X B[J, KX] → B[J, K]
     A[IX, J] · B[J, K] → C[IX, K]

   选哪个: 看后续算什么,通常由上下文决定

四条规则总结

  ┌─────────────────────────────────────────────────────┐
  │  Case 1: 都没切 contract 维       → 零通信           │
  │  Case 2: 一边切 contract 维       → AllGather 补齐   │
  │  Case 3: 两边同 axis 切 contract  → AllReduce 求和   │
  │  Case 4: 两边同 axis 切 non-contract → AllGather 一个│
  └─────────────────────────────────────────────────────┘

  通信代价对照:
    AllGather     = V / W
    ReduceScatter = V / W   (对称操作,代价一样)
    AllReduce     = 2V / W  (= RS + AG)
    AllToAll      = V / (4W) (单向环,比 AG 便宜 4 倍!)

这就是全书 Ch3 的核心。记住这 4 种 case,所有分布式 matmul 都变成"查表选 case"。

三、用这套记号重讲 Megatron TP

最爽的部分来了——我们用 [IX, JY] 记号重新讲一遍 Megatron TP,你会发现之前所有"直觉解释"都变成了严格的符号推导。

3.1 Column Parallel 其实是什么

Megatron 的 Column Parallel:权重 W ∈ [D, F] 按 F 维度（列）切开。

   用新记号写:
     Mesh({'X': 8})  (8 卡 TP)
     X: [B, D]        ← 输入,复制
     W: [D, FX]       ← 权重,F 维度切到 X
     Y = X · W: [B, D] · [D, FX] → ?

   套用 Case 1 判断:
     contract 维是 D
     X 的 D 没切,W 的 D 没切 → Case 1!
     ─────────────────────────────
     零通信!
     输出 Y[B, FX]  (结果自然 F 维度分片)

啊哈时刻:Column Parallel 的"零通信"不是直觉,是Case 1 的直接结果。

3.2 Row Parallel 其实是什么

Megatron 的 Row Parallel:权重 W ∈ [F, D] 按 F 维度（行）切开。

   接着上面:
     Y: [B, FX]        ← 输入已是 F 维度分片
     W2: [FX, D]       ← 权重 F 维度切 X
     Z = Y · W2: [B, FX] · [FX, D] → ?

   套用 Case 3 判断:
     contract 维是 F
     Y 的 F 切了 X,W2 的 F 切了 X,同一个 axis!
     ─────────────────────────────
     Case 3!先本地 matmul,再 AllReduce
     本地: Y · W2 → Z[B, D]{UX}
     通信: AllReduce_X → Z[B, D]

Column + Row 配对"天然免费衔接"的本质:

   Column 产出: Y[B, FX]     ← F 维度按 X 切
         ↓
         ↓ 中间状态天然就是 F 切 X
         ↓ 不需要任何通信!
         ↓
   Row 消费: 刚好需要 F 维度已经切了的输入

   如果反过来 Row + Column:
     Row 产出: Z[B, D]{UX} → AllReduce → Z[B, D]  (一次通信)
     Column 消费: 需要 Z[B, D] 完整的
                 (已经是了,因为 AllReduce 完就完整了)

     但是... 中间要存一个完整的 [B, D],比 [B, FX] 大 4-7 倍
     而且下一步还要重新切!

所以 Megatron 原论文的 Column + Row 配对,就是选了"中间状态天然衔接 + 只需一次 AllReduce"的最优路径。

3.3 Megatron 一层 Transformer 的通信

用新记号完整走一遍一层 MLP:

   输入: X[B, D]  (fully replicated,每卡都有完整 X)
   W_up: [D, FX]
   W_down: [FX, D]

   Step 1 (Column, Case 1): X · W_up
     [B, D] · [D, FX] → [B, FX]   零通信 ✓

   Step 2 (activation,元素级):  SiLU([B, FX]) → [B, FX]
     零通信 ✓

   Step 3 (Row, Case 3): [B, FX] · [FX, D]
     本地算: → [B, D]{UX}
     AllReduce: → [B, D]
     ─────────────────
     一次 AllReduce!

   一层 MLP 总通信:
     forward:  1 次 AllReduce
     backward: 1 次 AllReduce
     总共:    2 次

Attention 层也是同一个 pattern（Q/K/V 投影是 Column,output projection 是 Row）。一层 Transformer 总共 4 次 AllReduce（MLP + Attention,各 fwd/bwd 各一次）。

这个 4 从哪儿来:不是 Megatron 作者拍脑袋设计的,是Case 1 + Case 3 配对的数学结果。

四、用这套记号看 FSDP

现在用新记号重讲 FSDP,你会发现 FSDP = Case 2 的典型应用。

4.1 FSDP 的核心操作

   Mesh({'X': 8})   (8 卡 FSDP,也叫 DP=8 with sharding)

   初始状态:
     W: [D, FX]       ← 权重按 F 切到 X (每卡只有 F/8)
     B: [BX, D]       ← batch 按 X 切 (每卡只有不同 batch 数据)

   想做: Y = B · W

   套用 Case 2:
     contract 维是 D
     B 的 D 没切,W 的 D 没切
     但是... 等等,W 的输出维 F 切了 X,这不影响!
     再看一次: F 维度切 X,但那不是 contract 维

     修正:
       这个场景 contract 维确实都没切 → Case 1 吗?
       NO! 因为 batch 也切了 X

慢着,这里有个关键细节:FSDP 的实际操作不是直接做分片 matmul,而是先把参数 AllGather 回来再算。

4.2 FSDP 的真实时间线

   每层 forward 的真实流程:

   Step 0: 开始时 W 是分片状态 W[D, FX]
   Step 1: AllGather_X W[D, FX] → W[D, F]  ← 临时拉完整参数
   Step 2: B · W: [BX, D] · [D, F] → [BX, F]  ← Case 1,零通信
   Step 3: discard W,只保留本地那片

   backward 类似:
     再 AllGather 一次参数
     算 backward
     ReduceScatter 梯度回去

用 JAX 的话说,FSDP = "每层临时提升分片状态做计算,算完回落"。

4.3 FSDP 的通信代价

   每一步 iteration (fwd + bwd):

   fwd 每层:
     AllGather W: V/W  (V = 参数字节数)

   bwd 每层:
     AllGather W: V/W
     ReduceScatter grad: V/W

   总通信: 3V/W per layer

   对比纯 DP:
     只 AllReduce 梯度: 2V/W per layer
     ──────────────────
     FSDP = DP × 1.5x 通信

   但是:
     纯 DP 每卡存全部参数+梯度+优化器 = 16V (Adam)
     FSDP 每卡存 1/N = 2V (N=8 时)
     ──────────────────
     显存降到 1/8,通信增加 50%
     极值优: 适合参数量 >> batch 的大模型

面试锚点:FSDP 为什么能省显存?因为它用"AllGather + 立即 discard"的 trick,只在计算那一刻临时有完整参数。空间换时间的典型例子。

五、ReduceScatter 的深层意义

ReduceScatter 很多人只当成"AllReduce 的下半部分"。JAX Book 给了一个更深的视角:ReduceScatter 和 AllGather 是对偶的。

5.1 对偶性

  Forward 操作               Backward 对应
  ──────────────            ──────────────
  AllGather [IX] → [I]      ReduceScatter [I]{UX} → [IX]
  ReduceScatter [I]{UX} → [IX]   AllGather [IX] → [I]
  AllReduce [I]{UX} → [I]   AllReduce (对称)

直觉:AllGather 把分片"复制变完整";ReduceScatter 把"所有卡的 partial sum"合并再分片。它们是互为反函数的通信操作。

数学上:broadcast 和 reduce 是转置算子,AllGather 和 ReduceScatter 是它们的张量积,也互为转置。所以神经网络 forward 用 AllGather 就意味着 backward 自动用 ReduceScatter(链式法则)。

5.2 为什么有时候要用 ReduceScatter 代替 AllReduce

  场景: A[I, JX] · B[JX, K] → C
  默认: AllReduce → C[I, K]  (完整复制)
  优化: ReduceScatter → C[IX, K]  (保持分片!)
                   或   C[I, KX]  (选个维度分片)

  为什么好:
    AllReduce = 2V/W 通信
    ReduceScatter = V/W 通信  ← 省一半!

    代价: 结果仍是分片的
    如果后续操作本来就要分片输入 → 完美!
    如果后续非要完整的 → 延迟 AllGather 到真正需要时

这就是 Sequence Parallel (SP) 的核心:

  原 Megatron: Attention 输出 → AllReduce → LayerNorm 用完整输入
  Megatron + SP: Attention 输出 → ReduceScatter → LayerNorm 用分片输入
                                                 → 完成后 AllGather
  通信量不变! 但是 LayerNorm 的激活显存降为 1/N

面试锚点:SP 为什么通信量不变但省显存?因为 AllReduce = RS + AG,把它拆成两步并在中间插入 LayerNorm,激活值只在"RS 完成后 AG 前"这段分片存储。显存节省 N 倍,通信量完全不变。

六、AllToAll:MoE 的通信王牌

最后一个核心 collective:AllToAll。前面在分布式 Transformer里讲过 MoE 用它,这里补一个新视角——AllToAll = "搬一个下标"。

6.1 AllToAll 的记号

  AllToAll_{X,J}  A[IX, J] → A[I, JX]
                     ↑           ↑
                   X 切 I        X 切 J (改切了!)

  等价于: "把 X 这个 mesh axis 从 I 维度搬到 J 维度"

6.2 为什么 AllToAll 比 AllGather 便宜 4 倍

  AllGather 的做法:
    每张卡持有 V/N 字节,要把这些数据传给其他 N-1 张卡
    每张卡最终收到 (N-1)/N × V 字节

  AllToAll 的做法:
    每张卡持有 V/N 字节,但只需要传 V/N^2 给每张其他卡
    (因为每张卡发给另一张卡的数据量只是自己持有的 1/N)

  数学推导 (单向环):
    AllGather: 每卡发送 V/N × (N-1) ≈ V bytes
    AllToAll:  每卡发送 (V/N^2) × (N-1) × avg_hops
               ≈ V/N × 1/4
    (因为平均 hop 距离 = N/4 for 环形拓扑)

  结论: AllToAll ≈ AllGather / 4

MoE 为什么偏偏爱 AllToAll:MoE 的 token-to-expert 路由本质上就是"把 token 的 shard 维度从 'batch' 换到 'expert'"——完美对应"搬下标"的操作。

6.3 推理场景的实用场景

AllToAll 不只在训练里有用,推理也有一个经典用例:

  场景: DP Attention + EP FFN (现代 MoE 推理标准架构)

  Attention 阶段:
    每张卡独立处理不同请求 (DP)
    输出 h[B_X, D]   ← batch 按 X 切

  FFN 阶段 (MoE):
    需要按 expert 路由
    但现在 batch 切在 X 上,expert 切在...也在 X 上!

  AllToAll 救场:
    AllToAll_{X, expert} h[B_X, E] → h[B, E_X]
                                     ↑
                              batch 完整,expert 按 X 切

  每个 expert 独立计算,再 AllToAll 回来

面试锚点:MoE 推理为什么用 AllToAll 不用 AllReduce?因为 MoE 的本质是不同 token 走不同 expert,不是"所有卡算一个加法求和"。你需要的是"重新分组",不是"求和",所以 AllToAll 是最经济的操作。

七、Overlap 的艺术:让通信"消失"

讲到这儿你可能开始担心:每层都要 AllReduce / AllGather,这些通信加起来不就是巨大开销吗?

答案:现代框架都在尽量让通信和计算 overlap。

7.1 Collective Matmul 的思路

  传统做法 (串行):
    [────── matmul ──────][── AllReduce ──]
    GPU 算               GPU 等通信

  Collective Matmul (重叠):
    [── matmul tile 1 ──][── matmul tile 2 ──][── tile 3 ──]
                 [── AllReduce tile 1 ──][── AllReduce tile 2 ──]
                      ↑ 和下一个 tile 的计算重叠

  关键: 把大 matmul 切成多个 tile
       算完一个 tile 就开始它的 AllReduce
       同时开始下一个 tile 的计算

这就是 "collective matmul" 这个 kernel 名字的由来——把 collective 操作融合进 matmul 的 tiling 循环。

7.2 具体 overlap 的分类

  Overlap 类型 1: 不同层之间 (coarse-grained)
    例: FSDP 的 "prefetch"
    计算第 L 层时,后台 AllGather 第 L+1 层的参数

  Overlap 类型 2: 同一层内 tile 之间 (fine-grained)
    例: Megatron 的 "sequence parallel + context parallel 融合"
    一个大 matmul 切成多个 tile,通信和计算流水

  Overlap 类型 3: forward/backward 之间
    例: 1F1B pipeline schedule
    前面 micro-batch 的 backward 和后面的 forward 并发

  Overlap 类型 4: 不同通信操作之间
    例: DeepSeek V3 的 DualPipe
    把 AllToAll 塞进 pipeline bubble 里

面试锚点:为什么 TP 必须机内?因为 TP 的 AllReduce 太频繁,只有机内 NVLink 的带宽才能让通信时间 < 计算时间,否则 overlap 失败,通信变成瓶颈。

八、面试锚点

Q1:A[IX, JY] 这套记号比"Column/Row Parallel"好在哪?

A: 三点:(1) 可组合——A[IX, JY] 是独立下标,随便搭配;Column/Row 只有两种;(2) 可对任意 tensor 用——不局限于 2D 矩阵,3D/4D 张量都能描述;(3) 直接对应 collective——看一眼下标就知道该插哪个 collective(contract 维下标 → AllReduce,non-contract 维下标 → AllGather)。

Q2: 分布式 matmul 的四种 case 是哪四种?各自要哪个 collective?

A: Case 1(contract 维都没切,零通信); Case 2(一边切 contract 维,先 AllGather 补齐); Case 3(两边都切 contract 且在同一 axis,本地算完 AllReduce); Case 4(两边都切 non-contract 维在同一 axis,非法状态,先 AllGather 一个)。核心判断:contract 维有没有切。

Q3: AllGather / ReduceScatter / AllReduce / AllToAll 的 cost 关系?

A: 带宽上限区:AllGather = ReduceScatter = V/W, AllReduce = 2V/W (= AG + RS), AllToAll = V/(4W) (环形拓扑)。AllToAll 最便宜,AllReduce 最贵。延迟区(数据很小时)上述关系失效,全部变成 O(log N) 或 O(N) 跳数。

Q4: Megatron 的 Column + Row 配对为什么"天然免费"?

A: 用新记号看:Column Parallel 产出 Y[B, FX](F 维度自然按 X 切),刚好是 Row Parallel 的输入格式 [B, FX] · [FX, D]——contract 维两边都切了 X,直接套 Case 3 做本地 matmul + AllReduce。中间状态天然衔接,不需要额外通信。反过来 Row + Column 要先 AllReduce 再切,通信翻倍且中间显存更大。

Q5: FSDP 的通信代价为什么是 DP 的 1.5x?

A: 纯 DP 每 iteration 只做一次 AllReduce 梯度 = 2V/W。FSDP 每层:fwd 1 次 AllGather(V/W)+ bwd 1 次 AllGather(V/W)+ bwd 1 次 ReduceScatter(V/W)= 3V/W。总比值 3/2 = 1.5x。换来的是显存降到 1/N,大模型下 trade-off 极划算。

Q6: Sequence Parallel 为什么通信量不变但省显存?

A: 核心等式 AllReduce = ReduceScatter + AllGather。原 Megatron:Attention 输出 → AllReduce → LayerNorm。SP 改成:Attention 输出 → ReduceScatter → LayerNorm(用分片输入)→ AllGather。总通信量完全相同,但 LayerNorm 的激活值在"RS 之后 AG 之前"是分片存储的,显存降到 1/N。

Q7: ReduceScatter 和 AllGather 为什么互为对偶?

A: 从算子角度:broadcast 和 reduce 互为转置,AllGather 是 broadcast 的张量积,ReduceScatter 是 reduce 的张量积,所以它们也互为转置。从神经网络角度:forward 用 AllGather,backward 自动是 ReduceScatter(链式法则)。这就是为什么 FSDP forward 的 AllGather 参数对应 backward 的 ReduceScatter 梯度。

Q8: AllToAll 为什么比 AllGather 便宜 4 倍?

A: AllGather 每张卡要把自己的 shard 复制给其他所有卡(V/N × (N-1) ≈ V bytes);AllToAll 每张卡只需要给另一张卡发自己的 1/N(V/N²),加上平均 hop = N/4,总发送量约 V/(4N)。关键差异:AllGather 是广播,AllToAll 是重分发,后者天然利用了"不是所有数据都要去所有地方"。

Q9(进阶):MoE 推理为什么需要 AllToAll 而不是 AllReduce?

A: AllReduce 做的是"求和",前提是每张卡都在算同一个结果的部分和。但 MoE 本质是"不同 token 走不同 expert",需要的是重分组(把 batch 里属于 expert_i 的 token 送到持有 expert_i 的卡)。重分组对应"搬下标"操作,AllToAll 恰好是这个语义。强行用 AllReduce 需要大量 padding 和 masking,通信量和内存都爆炸。

Q10(进阶):给你一个新的并行方案,比如把 head 维度切到 X、hidden 维度切到 Y,怎么分析它?

A: 四步走:(1) 用 [IX, JY] 记号写出所有 tensor 的 sharding; (2) 对每一个 matmul,看 contract 维有没有被切,套 4 种 case 判断需要哪个 collective; (3) 累加每层的通信量,算 roofline(T_comm vs T_compute); (4) 看能不能 overlap(通信 < 计算则可完全隐藏)。整个流程纯机械,没有创造性,这就是记号系统的威力。

一句话记住这一篇

分布式 matmul 的所有复杂度,可以压缩成一句话:看 contract 维有没有被切,切了就要通信。[IX, JY] 记号让这个判断变成查表——Case 1 零通信,Case 2 用 AllGather,Case 3 用 AllReduce,Case 4 先 AllGather 掉一个。Megatron / FSDP / SP / MoE 不是四种独立技术,而是这套记号在不同 sharding 下的 4 种实例。学会了这套记号,未来看任何新并行方案都变成"套公式"。

系列下一篇预告:《Transformer Math:FLOPs/参数/KV Cache 精确算》——学会这篇后,你能对任何 Transformer 模型精确算出:(a) 参数量; (b) forward/backward FLOPs; (c) activation 显存; (d) KV cache 大小; (e) prefill 和 decode 的具体瓶颈。面试被问"Llama 70B 一个 batch 需要多少显存"再也不用猜。

配套阅读:JAX Scaling Book Ch3 原书章节(含交互式 JAX 代码)。相关文章:分布式 Transformer 深度解析、Roofline 分析、MoE 推理完整机制。

Sharded Matmul：分片矩阵乘的通用记号

Sharded Matmul：分片矩阵乘的通用记号

一、为什么需要一套"统一记号"

1.1 现有讨论的混乱

1.2 JAX Book 的解法:[IX, JY]

1.3 一个小例子感受一下

二、分片矩阵乘的 4 种情况

Case 1: 都没切 contract 维度 → 零通信

Case 2: 一边切了 contract 维度 → AllGather

Case 3: 两边都切 contract 维度在同一 axis → AllReduce

Case 4: 两边 non-contract 维度切在同一 axis → 先 AllGather

四条规则总结

三、用这套记号重讲 Megatron TP

3.1 Column Parallel 其实是什么

3.2 Row Parallel 其实是什么

3.3 Megatron 一层 Transformer 的通信

四、用这套记号看 FSDP

4.1 FSDP 的核心操作

4.2 FSDP 的真实时间线

4.3 FSDP 的通信代价

五、ReduceScatter 的深层意义

5.1 对偶性

5.2 为什么有时候要用 ReduceScatter 代替 AllReduce

六、AllToAll:MoE 的通信王牌

6.1 AllToAll 的记号

6.2 为什么 AllToAll 比 AllGather 便宜 4 倍

6.3 推理场景的实用场景

七、Overlap 的艺术:让通信"消失"

7.1 Collective Matmul 的思路

7.2 具体 overlap 的分类

八、面试锚点

一句话记住这一篇

评论

1.2 JAX Book 的解法:`[IX, JY]`