第9章：长序列训练与上下文并行

📖 本章概述

随着长文档、长上下文、长思维链需求的增长，序列长度从 2K 涨到 128K 甚至 1M。Attention 的计算量和激活显存与序列长度呈 $O(s^2)$ 关系，单卡根本放不下——这催生了沿序列维度切分的上下文并行（Context Parallelism, CP）。本章独立成章，讲清三种主流方案的切分维度差异。注意区分：上一章的序列并行（SP）切的是非 Attention 区域的激活，本章切的是 Attention 计算本身。

---

📑 章节结构

1. 序列长度墙

• Attention 的 $O(s^2)$ 计算与显存：$s$ 翻倍，开销翻 4 倍
• KV Cache 与激活显存随 $s$ 的增长
• 为什么 TP/PP/DP 都解决不了”单条序列太长”

2. FlashAttention 回顾（前置）

• 分块计算 + online softmax，把 $O(s^2)$ 显存降到 $O(s)$
• 为什么 CP 方案大多构建在 FlashAttention 之上

3. Ring Attention

• 沿序列维度把 Q/K/V 切到多卡
• 环形传递 KV 分片，每卡逐块计算局部 Attention，online softmax 合并
• 计算与 KV 传输重叠
• 因果掩码下的负载不均衡与 Zigzag/Striped 优化

4. DeepSpeed-Ulysses

• 沿 Head 维度切分 → All-to-All 重排为序列维度切分
• 通信量分析：与序列长度无关、与 Head 维度相关
• 与 Ring Attention 的取舍：通信模式 vs Head 数限制

5. Megatron Context Parallel

• 沿序列维度切分，配合 FlashAttention 分块
• 与 TP/SP 的协同：在 Megatron 中如何同时开 TP + CP
• 通信开销与序列长度的关系

6. 三种方案对比

方案	切分维度	通信原语	限制	适用场景
Ring Attention	序列	P2P 环传 KV	因果掩码负载不均	超长序列、KV 大
Ulysses	Head→序列	All-to-All	Head 数需 ≥ CP 度	Head 多的模型
Megatron CP	序列	P2P/AllGather	与 Megatron 绑定	Megatron 生态

7. 动手实验

• 用 FlashAttention 跑通单卡长序列，观察显存随 $s$ 的变化
• 实现简化版 Ring Attention 的 KV 环传 + online softmax 合并

---

🎯 本章学习目标

• 能解释 Attention 的 $O(s^2)$ 困境及为何需要沿序列切分
• 能区分 SP（非 Attention 激活）与 CP（Attention 计算）的切分维度
• 能对比 Ring Attention、Ulysses、Megatron CP 的切分维度与通信模式
• 能说明因果掩码给 Ring Attention 带来的负载均衡问题及优化思路