一层的矩阵乘,能不能拆到几张卡上一起算?
本课一问:ZeRO-3 之后,每张卡仍然一个人算完整的一层,并且存着自己那份样本的全部激活。能不能把一层里的矩阵乘拆开,几张卡一起算同一批 token?
这一课先不碰:按层切(第 4 课);注意力按头切只讲结论;长序列的切法(岔路 1)。
从已有出发
切什么每张卡一份完整模型,各算全局 batch 的 1/N;反向之后把 N 份梯度加起来,每张卡用同一份梯度更新。
多出的通信all-reduce 梯度
有多大锚点 BF16 梯度 55.6 GB;8 卡 ring all-reduce 每卡每步发 97.2 GB。
切什么优化器状态、梯度、参数按 N 份切开,每张卡只存、只更新自己那份;ZeRO-3 在用到一层之前把这一层的参数拼回来。
多出的通信reduce-scatter 梯度 + all-gather 参数
有多大8 卡 ZeRO-3:每卡模型状态 55.6 GB(原来 444.5 GB);通信每卡每步 145.9 GB,是 DP 的 1.5 倍。
切什么一层里的矩阵乘拆给 t 张卡:第一个矩阵按列切、第二个按行切,非线性各算各的,最后一次 all-reduce。
多出的通信每层 all-reduce(前向 2 次、反向 2 次)
有多大t = 8、每 4,096 个 token:每卡发 18.8 GB,每层都要等它做完;所以只放在一台机的 8 张卡里。
切什么把 64 层按顺序切成 p 段放在不同的卡上,每段做完把激活交给下一段,反向时再把梯度传回来。
多出的通信段与段之间 send / recv 激活
有多大锚点每段 8 层(p = 8);每个 micro-batch(4,096 个 token)在段间传 0.0 GB 激活;代价是气泡:p = 8、m = 32 时 21.9%。
切什么一条很长的序列切给几张卡,每张卡只存自己那段 token 的激活;注意力需要全部 K、V,就让 K、V 块沿环传一圈。
多出的通信注意力里沿环传 K、V
有多大s = 131,072、8 卡:每张卡每一步传出一块 K、V 0.1 GB,每个全注意力层传 7 块。
切什么MoE 的专家分到不同卡上,token 用 all-to-all 发到持有对应专家的卡,算完再 all-to-all 收回来。
多出的通信FFN 前后各一次 all-to-all
有多大DeepSeek-V3:256 个路由专家 ÷ 64 张卡 = 每卡每层 4 个;每个 token 最多去 4 个节点。
上一课结束时的账:8 卡 ZeRO-3,每卡模型状态 55.6 GB,剩 24.4 GB 给激活;s = 4,096 时 C9 的选择性重计算要 45.6 GB,放不下。原因是 ZeRO 仍是数据并行:每张卡处理自己的样本、算完整的每一层,激活一点没切。
这一课换个切法:8 张卡处理同一批 token,每张卡只拿一层权重的一部分。贯穿例句「小明把苹果给了小王」回来了——它的 5 个 token 在 C3 教学世界里是一个 5 × 4 的矩阵 X,下面的验算台就用它当输入:每张卡拿到同一个 X、不同的一块权重。
先自己设计:把第一个矩阵沿收缩维切开
FFN 是两次矩阵乘夹一个非线性:Y = f(X·A)·B。X·A 的每个元素是 X 的一行和 A 的一列逐项相乘再求和——最直接的想法是把这个求和拆开:卡 0 算前一半维度的乘积和,卡 1 算后一半,各自做完非线性和第二次乘法,最后加起来。
- ① 朴素方案
- ② 它在哪失败
- ③ 改进方案
试三件事:① 在朴素一侧看红色格子的最大差,再点「换一组权重」——每一组都错,不是这组数碰巧;② 在改进一侧依次切三种切法,对照「这种切法需要的通信」:列切 → 行切只要最后一次 all-reduce,按行切要在非线性之前多一次大得多的同步,两个都按列切要 all-gather 整个中间结果;③ 把卡数切到 4、非线性切到 GeLU:列切 → 行切仍然逐元素相等(剩下的差是 10⁻¹⁶ 量级的浮点舍入)。
改进:先按列切,再按行切
Megatron-LM(Shoeybi 等,2019)的切法只用到分块矩阵乘法:
这就是 Megatron-LM 原文的说法:第一个 GEMM 按列切、第二个按行切,每层只需要前向两次、反向两次 all-reduce 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条分块矩阵乘的恒等式加「非线性逐元素」,对任意权重成立;原文写的是「前向两次、反向两次 all-reduce」。 出处Megatron-LM §3(需联网) ↗需联网 。放回锚点 Qwen3.8 的真实形状(t = 8):
| 维度 | 含义 | 真实规模 | 小规模 |
|---|---|---|---|
| T | 这批 token 的数量(贯穿例句是 5 个) | 5 | 5 |
| d | 隐藏维度 hidden_size | 5120 | 4 |
| ft | 每张卡分到的中间维 = intermediate_size ÷ t(t = 8;小规模对应验算台的 16 ÷ 2) | 2176 | 8 |
锚点的注意力层要多看一眼:每层 24 个 Q 头,t = 8 时每张卡 3 个;可 KV 头只有 4 个,t = 8 时不够每张卡一个,每个 KV 头要复制到 2 张卡上(t = 4 时正好每卡一个)。GQA 省下的 KV,在张量并行里变成了「切不开」——这是一个由 config 直接决定的约束。
代价:每一层都要等
张量并行的通信和数据并行完全不同。数据并行每步同步一次梯度,量由参数量决定;张量并行每一层都要同步激活,量由 token 数决定。按锚点算,t = 8、一次处理 4,096 个 token:每次 all-reduce 的张量是 4,096 × 5,120 × 2 字节,ring 下每卡发 73.4 MB;一层前向 2 次、反向 2 次,64 层就是每卡 18.8 GB——而这 4,096 个 token 的计算,摊到每张卡只有 6 × 25,624,600,064 × 4,096 ÷ 8 ≈ 79 TFLOP(只数参与矩阵乘的参数:去掉只查不乘的 embedding、视觉塔与 MTP,与 C11 第 1 课同一口径)。
这些同步夹在层与层之间,下一层必须等它做完才能开始,很难像数据并行那样藏到反向后面。所以带宽决定一切:按 DeepSeek-V3 报告里 H800 集群的数字,节点内 NVLink 160 GB/s、节点间 IB 50 GB/s 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条报告原话:NVLink 提供 160 GB/s,约是 IB(50 GB/s)的 3.2 倍。这是 H800 集群的数;H100 的 NVLink 规格更高。 出处DeepSeek-V3 报告 §3.2.2(本地 PDF) ↗ ,18.8 GB 在机内要约 117 ms,跨机要约 376 ms(推演,不含重叠)。实测更直接:ZeRO 论文测过把 40B 模型用 Megatron 的张量并行跨两台 DGX-2,每张 V100 只有约 5 TFLOPS,不到峰值的 5% 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条原文:「about 5 Tflops per V100 GPU (less than 5% of hardware peak)」。是他们的实测。 出处ZeRO 原文 §1(需联网) ↗需联网 。于是有了现代配置里的一条共识:张量并行放在最内层,只在一台机的 8 张卡里用。Llama 3 的 4D 并行里 TP = 8 在最内层;DeepSeek-V3 更进一步,完全不用张量并行,靠显存优化规避了它 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条DeepSeek-V3 原话是「without using costly Tensor Parallelism」;Llama 3 表 4 三种配置的 TP 都是 8。 出处DeepSeek-V3 §3.2;Llama 3 表 4 ↗ 。
回报:激活也切开了
张量并行不只切参数,也切激活:列切之后,每张卡只存 H 的 1/t。再配上 Korthikanti 等人的序列并行——把 LayerNorm、dropout 这些按 token 独立、原本在每张卡上重复计算的部分按序列切开,all-reduce 换成等量的 reduce-scatter + all-gather——每层激活整体变成原来的 1/t 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条激活公式 sbh(34/t + 5as/(ht)) 是逐张量记账;论文报告序列并行 + 选择性重计算让激活降到 1/5,530B 模型 MFU 从 42.1% 提到 54.2%(实测)。 出处Korthikanti et al. 2022(需联网) ↗需联网 。锚点 s = 4,096:选择性重计算的 45.6 GB 在 t = 8 时降到 5.7 GB。这正是上一课 ZeRO 做不到的事。
常见误解之一:「张量并行就是把参数分开存。」
参数分开存的是 ZeRO(用时再拼回来)。张量并行是把计算分开:每张卡永远只算自己那几列、那几行,处理的是同一批 token,每一层都要同步一次激活。
常见误解之二:「矩阵怎么切都行,反正最后加起来。」
切错维度,非线性之前就得多一次同步(上面的朴素做法),或者要 all-gather 整个 17,408 宽的中间结果——是 all-reduce 5,120 宽输出的 3.4 倍字节。「先列后行」是让非线性落在不需要通信的地方。
常见误解之三:「TP 越大越好。」
每层 4 次同步,量跟着 token 走;跨机就被 IB 拖垮。锚点只有 4 个 KV 头,t = 8 时还要复制。现代配置里 TP 通常不超过 8。
🔍 反向的 2 次 all-reduce 从哪来?(Megatron 的 f 与 g)
Megatron 原文把通信写成两个互为共轭的算子:f 放在列切矩阵之前,前向什么也不做(X 本来就每卡一份),反向要把各卡对 X 的梯度 all-reduce(每张卡只算了自己那几列对 X 的贡献);g 放在行切矩阵之后,前向 all-reduce(把部分和加起来),反向什么也不做。一个 MLP、一个注意力块各有一对 f、g,所以一层前向 2 次、反向 2 次。序列并行把 g 的 all-reduce 换成 reduce-scatter、把 f 换成 all-gather,通信量不变,激活却按序列切开了。
证据三问
| 问 | 本课的回答 |
|---|---|
| ① 它解决什么问题? | 一层的计算和激活太大:让 t 张卡处理同一批 token、各算一层的 1/t |
| ② 数学上为什么可能有效? | 分块矩阵乘 + 非线性逐元素:先列后行,切分前后结果逐元素相等 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条对任意权重成立;验算台里的剩余差是浮点加法换了结合顺序的舍入。 出处Megatron-LM §3 ↗需联网 ;激活按 1/t 缩小(配序列并行) 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条逐张量记账。 出处Korthikanti 2022 ↗需联网 |
| ③ 实验上为什么最后选了它? | Megatron-LM 用 512 张 GPU 训 8.3B 模型,达到 15.1 PFLOPs、76% 的扩展效率;Llama 3 的 TP 是 8;跨机时 ZeRO 原文实测只有峰值的 5% 以下,DeepSeek-V3 干脆不用 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条都是作者自报的实测或配置。 出处Megatron-LM 摘要;ZeRO §1;DeepSeek-V3 §3.2 ↗需联网 |
演进地图:你在这里
- 解决了训练一步的全部东西都在一张卡上,不需要任何通信。卡在锚点的模型状态 16 字节 × 27,781,427,952 = 444.5 GB,是 80 GB 的 5.6 倍;就算放得下也太慢:DeepSeek-V3 预训练用了 2,664K H800 GPU 小时,一张卡约 304 年。
- 解决了每卡一份完整模型、算 1/N 的数据,每步把梯度加起来。ring 让每卡只收发 2(N−1)/N 倍梯度:锚点 8 卡 97.2 GB、64 卡 109.4 GB(汇总到 0 号卡的朴素做法:388.9 GB → 3,500 GB)。卡在只解决「慢」:每张卡仍要放下完整的 444.5 GB。
- 解决了反向逐层产出梯度,攒满一桶(默认 25 MiB)就先发,通信藏到后面几层的反向里。卡在DP 的通信量每步固定:锚点 8 卡每卡 97.2 GB,按 IB 50 GB/s 要 1.94 s;每卡 4,096 个 token 的反向只有约 1.05 s(推演)——token 太少就藏不住。
- 解决了不再每卡存一份:8 卡上 444.5 → 152.8 → 104.2 → 55.6 GB(依次切优化器状态、梯度、参数)。卡在切参数后每步多一次 all-gather,通信是 DP 的 1.5 倍;激活不切:s = 4,096 的选择性重计算要 45.6 GB,放不进剩下的 24.4 GB。
- 解决了第一个矩阵按列切、第二个按行切,一个 MLP 只要一次 all-reduce;t = 8 配序列并行,激活 45.6 → 5.7 GB。卡在每层前向 2 次、反向 2 次同步:每 4,096 个 token 每卡要发 18.8 GB,每层都要等——只能放在机内。ZeRO 原文实测:跨两台 DGX-2,每卡约 5 TFLOPS,不到峰值 5%。
- 解决了按层切,跨机只传段与段之间的激活;切 micro-batch 填满流水线;1F1B 把在途激活从 m 份压到 p 份;interleaved 把气泡除以 v。卡在首尾必然空等:p = 8、m = 32 时气泡 (p−1)/m = 21.9%(占总时间 17.9%),interleaved v = 4 也还有 5.5%。
- 解决了反向拆成「对输入」B 与「对权重」W,W 拿去填空泡(F = B = W 时 ZB-H1 的气泡是 1F1B 的 1/3);DualPipe 从两头喂 micro-batch,并把 all-to-all 藏到另一批计算后面。卡在DualPipe 要存两份参数;调度更复杂;通信只是被藏住了,量一点没少。
- 解决了128K 的序列按序列切给多张卡;zig-zag 头尾配对,把因果负载 10 : 26 : 42 : 58 配平成每卡一样多。卡在每个全注意力层都要在卡之间传 K、V:锚点 s = 131,072、8 卡时每层 0.5 GB;Ulysses 按头切,KV 只有 4 个头。
- 解决了DeepSeek-V3 的 654B 路由专家参数分到 64 张卡,每卡每层 4 个;每个 token 最多去 4 个节点。卡在跨节点 IB 50 GB/s 只有 NVLink 160 GB/s 的约 1/3,跨节点 EP 的计算与通信约 1 : 1;热门专家所在的卡成为掉队者(→ C11)。
- 解决了Llama 3 405B:[TP, CP, PP, DP] = 8 × 1 × 16 × 128,16,384 张 H100,BF16 MFU 41%;DeepSeek-V3:PP 16 × EP 64 × ZeRO-1 DP,不用 TP。卡在没有万能配置:每种组合都要先算清每卡显存、每一维的通信走哪条链路、气泡多大(实验课)。
一句话
张量并行把一层的矩阵乘拆给 t 张卡:第一个矩阵按列切、第二个按行切,非线性在中间各算各的,每个 MLP 前向只需一次 all-reduce;通信跟着 token 数走、每层都要等,所以只放在节点内。
先停在这里:记住「先列后行」「每层前向 2 次、反向 2 次同步」「TP 只放机内」三句,就可以进下一课。
再深一层:展开 f 与 g 的 Dig,再在验算台里把「两个都按列切」和「列切 → 行切」的通信并排比一比。
下一个缺口
TP 只能在一台机的 8 张卡里用;卡再多,跨机该按层切——按层切为什么有卡在等? → 第 4 课:按层切到不同卡上,为什么有卡在等?
你可能会问
- 「张量并行和 ZeRO-3 都把参数切开了,区别到底在哪?」 ZeRO-3 的卡处理不同的样本,用到一层时把完整参数拼回来、自己算完整的一层;张量并行的卡处理同一批 token,参数永远不拼,每张卡只算自己那一块,同步的是激活。前者通信量看参数、后者看 token。
- 「为什么注意力按头切不需要中间通信?」 多头注意力里每个头的 Q、K、V 只和自己那个头有关(C3 第 7 课),头与头之间直到输出投影才汇合;按头切 Q、K、V,再按行切输出投影,和 MLP 的「先列后行」是同一个结构。
- 「DeepSeek-V3 不用 TP,那一层的激活怎么放下?」 报告 §3.2.3 列了它的做法:反向时重算所有 RMSNorm 与 MLA 的上投影、把指数滑动平均放到 CPU、让 MTP 与主模型共享 embedding 和输出头;再加上 16 路流水线把层切开、64 路专家并行把专家切开。