长上下文序列并行训练深度解析:从 Ring Attention 到 Ulysses 的百万 token 分布式注意力架构
长上下文序列并行训练深度解析:从 Ring Attention 到 Ulysses 的百万 token 分布式注意力架构核心痛点:当上下文从 32K 扩展到百万 token,单卡显存被注意力激活与 KV 占满、单步计算量爆炸,而数据并行无法切分序列维度、朴素张量并行又受限于头数与互联带宽——长上下文训练卡在"序列维不可拆"这一道墙上适配人群:大模型预训练/继续训练工程师、分布式训练框架开发者、关注百万级上下文(Llama 3 128K、Qwen 1M、长文档/长代码)训练的算法与系统工程师、有 Megatron-LM 或 DeepSpeed 经验并想深入序列并行内核机制的开发者收获能力:掌握序列并行(Context Parallelism)的统一视角与在线 softmax 可聚合性、DeepSpeed-Ulysses 全互联换维机制及其头数上限、Ring Attention 环形 K/V 传递与计算通信重叠及其通信不随卡数下降的瓶颈、Striped/zigzag 负载均衡重排、USP 二维混合如何同时突破头数上限与环通信、以及生产部署中 CP 与 TP 头维冲突、因果掩码负载不均、decode 不走 CP、变长序列 DCP 等工程陷阱与应对技术背景与演进逻辑背景:自回归与长上下文让序列长度 S 从 2K 一路涨到 128K 甚至 1M,注意力成为显存与算力的双重瓶颈算力:每层注意力计算量正比于S 2 S^2S2,S 翻百倍则单层算力翻万倍显存:朴素实现需物化S × S S × SS×S注意力分数矩阵,百万 token 下单层即数十 GBKV 代价(块公式):M m a t h r m K V = 2 c d o t L c d o t d m a t h r m m o d e l c d o t S c d o t B c d o t b M_{mathrm{KV}} = 2 cdot L cdot d_{mathrm{model}} cdot S cdot B cdot bMmathrmKV=2cdotLcdotdmathrmmodelcdotScdotBcdotb(2 表示 K 与 V,L 为层数,d_model 为隐藏维,S 为序列长,B 为批,b 为每元素字节数)推演:单卡无论算力还是显存都无法承载百万 token 的完整注意力 - 必须把序列维 S 拆到多卡线索一:数据并行切不动序列现象:数据并行按 batch 切,每卡仍要处理完整序列,单卡注意力依然是O ( S 2 ) O(S^2)O(S2)推论:长上下文的瓶颈在序列维而非 batch 维 - 需要一种沿序列切分的并行线索二:张量并行切不动规模现象:张量并行按头切注意力,能分摊算力,但 CP 度受限于注意力头数 H,且与 Ulysses 争抢头维推论:纯 TP 无法把序列切到几十上百卡 - 需要一种不依赖头数、专门切序列的并行线索三:FlashAttention 让注意力可分块聚合机制:在线 softmax(online softmax)逐块累积局部结果,无需物化S × S S × SS×S矩阵推论:既然注意力可按 K/V 块增量聚合,那么把不同块放到不同卡、再通信聚合,就能得到与单卡等价的精确注意力 - 分布式注意力在数学上成立早期序列并行的局限代表:2021 年 Sequence Parallelism(Li 等)与 DistFlashAttn缺陷一:仍按传统流程先算Q K m a t h r m T Q K^{mathrm{T}}QKmathrmT再掩码,需要S × S S × SS×S临时缓冲,长序列下显存炸裂缺陷二:与张量并行强绑定(把 all-reduce 拆成 all-gather 加 reduce-scatter),无法独立使用结论:现代长上下文需求倒逼出"FlashAttention 兼容、可独立使用"的新一代序列并行演进时间线(知识表格):年份里程碑核心贡献2022.06FlashAttention在线 softmax,免物化 S×S 矩阵,奠定可分块聚合基础2023.05DeepSpeed-Ulysses全互联换维,token 维切到头维,通信低但受头数限制2023.10Ring Attention环形传 K/V 加块级 FlashAttention,突破头数限制2023.11Striped Attention条带重排,解决因果掩码下 Ring 的负载不均2024.05USP 统一序列并行Ulysses 维乘 Ring 维二维混合,兼得两者优势2024.06LoongTrain头-上下文并行,另一种二维混合实现2024.07Llama 3 报告Meta 在 128K 训练中正式采用 Context Parallelism2025.05RingX近线性扩展至 128 卡,缓解 Ring 通信瓶颈2025.10DCP动态上下文并行,解决变长序列输入的负载抖动2026.01yunchang 0.6.4USP 工业实现支持 FA3/FA2/NPU 与 zigzag/stripe 重排总结:技术迭代的必然性在于"注意力在线 softmax 可分块聚合"提供了把序列拆到多卡仍精确的数学前提,核心价值是让百万 token 的注意力算力与显存在卡间被真正摊薄,而非靠单卡硬扛核心原理深度解析序列并行的统一视角