128K 的序列,激活放在哪?
这是一堂岔路课。主线的四种切法——数据、模型状态、层内矩阵、层——都默认「一条序列放得进一张卡」。长上下文训练时这一条不成立了。
本课一问:一条 128K 个 token 的序列,激活放在哪?
这一课先不碰:长上下文的位置编码(C2 第 5 课)、训练阶段怎么把长度一步步拉长(C11)、推理时的 KV cache(C3、C13)。
从已有出发
切什么每张卡一份完整模型,各算全局 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 个节点。
C9 第 5 课的激活公式里,序列长度 s 出现两次:线性项 34sbh 和注意力分数项 5as²b。FlashAttention 不把 s × s 的分数表写进显存,全量重计算每层只存输入 2sbh——这两招把激活压到了头,可它们都随 s 线性增长。这一课讨论长序列,贯穿例句只有 5 个 token,太短,下面的组件改用位置编号 0、1、2……
先自己设计:整条序列放一张卡
- ① 朴素方案
- ② 它在哪失败
- ③ 改进方案
| 锚点 Qwen3.8-27B · 一条序列 · BF16 | s = 32,768 | s = 131,072 | s = 262,144(原生) |
|---|---|---|---|
| 物化的注意力分数(单头单层,s × s × 2 字节) | 2.1 GB | 34.4 GB | 137.4 GB |
| 全量重计算:每层只存输入 2sbh × 64 层 | 21.5 GB | 85.9 GB | 171.8 GB |
数字口径:按 Korthikanti 2022 的记账与锚点 config 现算,b = 1,GB = 10⁹ 字节。还没算模型状态。
切给 8 张卡,每张卡只存自己那 16,384 个 token 的激活:每层输入 2sbh 那一项降到 10.7 GB。难点在注意力:第 i 个 token 的 query 要看它前面所有 token 的 key 和 value,而它们散在别的卡上。下面三种做法都在解决这件事。
改进一:Megatron 的序列并行(配合张量并行)
第 3 课的张量并行把注意力和 MLP 切开了,但 LayerNorm、dropout 这些按 token 独立的部分,在每张卡上各算一整份。Korthikanti 等人(2022)把这些部分按序列切开,原来的 all-reduce 换成等量的 reduce-scatter + all-gather,激活整体变成 1/t 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条逐张量记账;论文实测激活降到 1/5、重计算开销减少 90% 以上。 出处Korthikanti et al. 2022(需联网) ↗需联网 。它不增加通信,但只在张量并行的组内切,切的份数 = t ≤ 8。
改进二:Ring Attention / 上下文并行,K、V 沿环走一圈
Q 留在本卡,把 K、V 按序列切成 N 块,沿环传 N − 1 步:第 k 步,卡 c 拿到卡 c − k 的那一块 K、V,算自己的 query 对这块 key 的注意力,在线累加(和 FlashAttention 的分块在线 softmax 是同一个技巧),再把这块传给下一张卡。传的同时算上一块,通信可以完全藏在计算后面 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条摘要:分块计算,并把 K、V 块的通信与分块注意力的计算完全重叠。更早的 Li et al. 2021(arXiv 2105.13120)提出了 Ring Self-Attention。 出处Liu et al. 2023(Ring Attention,需联网) ↗需联网 。
但因果注意力有个麻烦:第 i 个 query 只看 j ≤ i 的 key。把序列连续切给 N 张卡,卡 0 的 query 只能看自己那一段,最后一张卡要看全部:
试三件事:① 读条形图:连续切分时 4 张卡要算的 (i, j) 对数是 10 : 26 : 42 : 58,最慢的卡是平均的 1.71 倍;② 看「Ring 的 N 步」表:每一步都要等最慢的那张卡,很多卡在一些步里算 0 个对,总耗时 58,而完全均衡只要 34;③ 切到 zig-zag,每张卡都是 34,Ring 的总耗时也降到 34——再把示例长度换成 32、卡数换成 8,仍然精确相等。
zig-zag 的做法:把序列切成 2N 段,卡 k 拿第 k 段(靠前、要算的少)和第 2N − 1 − k 段(靠后、要算的多)。第 k 段有 k·a² + a(a + 1)/2 个对(a 是每段长度),k 与 2N − 1 − k 两段相加,结果与 k 无关——每张卡的工作量按构造相等 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条k·a² + (2N−1−k)·a² + a(a+1) = (2N−1)·a² + a(a+1),与 k 无关。组件对每种 N、s 逐 token 计数核对。 出处待补 。Llama 3 的上下文并行用的正是这种配对:把序列切成 2 × CP 段,第 i 张卡拿第 i 段和第 2 × CP − 1 − i 段 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条报告原文的切法;它的上下文并行不走 Ring,而是先 all-gather 全部 K、V 再算,理由是更灵活、K、V 比 Q 小得多(GQA)。 出处Llama 3 报告 §3.3.2(需联网) ↗需联网 ——不过它不沿环传,而是先把 K、V all-gather 到每张卡上再算。
改进三:Ulysses,换成按头切
DeepSpeed-Ulysses(2023)换了个角度:进注意力之前做一次 all-to-all,把「每张卡有一段 token 的全部头」换成「每张卡有全部 token 的一部分头」;注意力算完,再 all-to-all 换回来 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条摘要:序列长度与设备数同比例增长时,通信量保持不变。 出处Jacobs et al. 2023(DeepSpeed-Ulysses,需联网) ↗需联网 。多头注意力的头之间互不相干,所以按头切的注意力不需要任何中间通信——和第 3 课按头切的张量并行是同一个道理。
它的约束来自 config:头数必须能被卡数整除。锚点有 24 个 Q 头,8 张卡正好每卡 3 个;可 KV 头只有 4 个,切给 8 张卡就得复制。Ring 没有这个限制,但每个全注意力层都要传 K、V 块。锚点 s = 131,072、8 张卡时,Ring 每个全注意力层每卡传 0.5 GB,Ulysses 的四次 all-to-all(Q、K、V、输出)每卡传 0.4 GB——量级相同,约束不同。下面把这些数放在一起:
线性注意力层:只传一个固定大小的状态
锚点 64 层里只有 16 层是全注意力;另外 48 层是线性注意力(Gated DeltaNet,C11 揭晓)。线性注意力把前面所有 token 压成一个固定大小的状态,锚点每层是 48 × 128 × 128 = 786,432 个数,和序列长度无关。做上下文并行时,卡与卡之间只需要接力传这个状态,而不是随 s 增长的 K、V 块。
Kimi K3 的报告把这一点写成了一节:softmax 注意力做上下文并行要交换随序列长度增长的 K、V 块,线性注意力只带一个固定大小的状态;但 delta rule 这类更新(K3 的 KDA、Qwen3.8 的 GDN 都属于这一类)会先把进来的状态乘一个与 token 有关的矩阵,不能把各段的局部状态直接相加,K3 为此多传一个「累计转移」,只需一次固定大小的 all-gather 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条报告的描述与实现(KCP);K3 预训练同时用了流水线(含虚拟段)、专家并行、ZeRO-1 数据并行与上下文并行。Qwen3.8 的训练并行方式未公开。 出处Kimi K3 报告 §5.1.2、§5.2(本地 PDF) ↗ 。这是 2026 年混合架构的一个隐性好处:长序列的通信只由那少数几层全注意力决定。
常见误解之一:「序列并行和张量并行是两回事,不能一起用。」
Megatron 的序列并行本来就是张量并行的配套:在张量并行的组内,把 LayerNorm、dropout 这些部分按序列切开,通信量不变。
常见误解之二:「把序列切成等长的块,负载就均衡了。」
因果掩码让靠后的 token 看得更多:4 卡连续切分是 10 : 26 : 42 : 58,序列很长时趋近 1 : 3 : 5 : 7。要 zig-zag 这类配对切法才均衡。
常见误解之三:「上下文并行就是把 KV cache 分到几张卡上。」
那是推理时的事。训练时切的是每一层的激活和计算,K、V 块在前向、反向里都要传。
证据三问
| 问 | 本课的回答 |
|---|---|
| ① 它解决什么问题? | 一条序列太长,激活放不进一张卡:按序列维切给几张卡 |
| ② 数学上为什么可能有效? | 注意力可以按 key 块分块计算、在线累加(与 FlashAttention 同理);zig-zag 的负载均衡按构造成立 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条分块在线 softmax 的结果与一次算完完全相同;zig-zag 的等式见上文。 出处待补 |
| ③ 实验上为什么最后选了它? | Ring Attention 与 Ulysses 都报告了随卡数扩展的长序列训练;Llama 3 在 131,072 长度的阶段用了 16 路上下文并行(表 4);K3 为线性注意力层单独做了上下文并行 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条作者自报的配置与实测。 出处Llama 3 表 4;Ring Attention;Ulysses;K3 §5.1.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。卡在没有万能配置:每种组合都要先算清每卡显存、每一维的通信走哪条链路、气泡多大(实验课)。
一句话
长序列的激活按序列切到多张卡上;因果注意力让连续切出的块工作量不等,zig-zag 把头尾配对就精确均衡;Ring 沿环传 K、V 块,Ulysses 用 all-to-all 换成按头切,后者受头数限制。
回到主线
实验课:几种切法怎么组合成一套方案? → 实验课(设备网格里有 CP 这一维);或者先看另一条岔路:MoE 的专家放在哪?
你可能会问
- 「zig-zag 会打乱 token 的顺序吗?」 不会改变模型看到的顺序:每个 token 的位置编号不变,因果掩码按原始位置算,只是「哪张卡负责哪几段」换了。
- 「Ring 和 Ulysses 能一起用吗?」 能。USP(Fang 与 Zhao,2024,arXiv 2405.07719)把两者组合成二维的序列并行:一个维度按头切(Ulysses),另一个维度沿环传(Ring),绕开 Ulysses 的头数限制。本章不展开。
- 「长上下文训练为什么通常放在最后阶段?」 激活与注意力的算量都随 s 增长,前期用短序列更划算,最后再把长度拉长——C11 讲长上下文训练时会回到这里。