为什么长上下文会先撞上激活显存
长上下文训练的难点不只是 Attention 计算量增长。即使使用 FlashAttention 降低中间矩阵的物化开销，Transformer 各层仍需要保留与序列长度近似线性增长的激活，用于反向传播。序列从 8K 扩展到 128K 后，模型参数没有变化，但单卡保存的激活、位置编码、Mask 和 Attention 工作区都会迅速挤占显存。
常见应对方法有三种：
| 方法 | 优势 | 代价 |
|---|---|---|
| 减小 Micro Batch | 直接降低显存 | 矩阵计算效率下降 |
| 增加 Activation Checkpointing | 用计算换显存 | 重计算增加算力开销 |
| 提高 Tensor Parallel 度数 | 分摊 Linear 层显存 | 通信更难被计算覆盖 |
Context Parallel（CP） 换了一个切分维度：不再让每张 GPU 保存完整序列，而是沿序列维度切分输入和激活。若 CP Size 为 4，一段长序列会被分成四个 Context Shard，每张 GPU 只处理其中一部分 Token。这样，逐 Token 的 Linear、MLP、LayerNorm 等算子基本可以在本地执行，激活显存也随分片数下降。
真正困难的部分只有 Attention：本地 Query 仍然需要看到全序列的 Key 和 Value。
核心原理：切分 Q，流转 KV
设序列长度为 S，Context Parallel 度数为 C。每个 Rank 保存约 S/C 个 Token 的 Q、K、V。对本地 Q 来说，Attention 输出仍要求遍历完整 K/V，因此系统必须在多个 Rank 之间交换 KV Shard，并把各分块的 Softmax 统计量正确合并。
这不是简单地把多个局部 Attention 结果做平均。在线 Softmax 需要维护每一行的最大值和指数和，才能在不物化完整注意力矩阵的情况下，将不同 KV Block 的结果精确合并。
Ring Attention：让 KV 沿环移动
Ring Attention 保留本地 Q，让每个 Rank 的 KV Block 沿环形拓扑逐轮传递：
- 使用本地 Q 与当前 KV Block 计算局部 Attention；
- KV 异步发送给下一个 Rank，同时接收上一个 Rank 的 KV；
- 更新在线 Softmax 状态；
- 重复直到本地 Q 看过所有 KV Block。
关键不是”用了 Ring”本身，而是能否让下一轮 KV 通信与当前轮 Attention Kernel 重叠。若 Block 太小，Kernel 时间不足以隐藏通信；若 Block 太大，显存峰值、流水线气泡和尾部延迟又会上升。
Ulysses：用 All-to-All 重排序列和注意力头
DeepSpeed-Ulysses 采用另一条路线。输入先按序列维度切分，进入 Attention 前通过 All-to-All 将数据从”每卡持有部分序列、全部或部分头”重排为”每卡持有完整序列的一部分注意力头”，完成 Attention 后再通过 All-to-All 还原。
它的优势是通信模式清晰，且在合适拓扑上容易获得高带宽；工程限制则通常来自注意力头数量、GQA/MQA 的 KV Head 数、分片可整除性，以及跨节点 All-to-All 的尾部延迟。不能只看 GPU 数量决定 Ulysses Degree，必须同时检查 Head Layout。
All-Gather、P2P Ring 与 All-to-All 不是固定答案
PyTorch Context Parallel API 同时提供基于 All-Gather 的 Pass-KV 和基于 All-to-All 的轮转方式；Megatron Core 则会在底层将 KV 交换转换为环形点对点通信，并利用 MQA/GQA 较少的 KV Head 降低通信量。
生产选型应做拓扑实测：
- 单机 NVLink/NVSwitch 环境：All-to-All 或集合通信通常更容易获得稳定带宽；
- 跨节点：需要比较 InfiniBand/RoCE 的 P2P、All-Gather 和 All-to-All 表现；
- 不能只看平均带宽，还要看最慢 Rank、通信启动次数和计算—通信重叠比例；
- Attention Backend、NCCL 版本、CUDA Graph 与动态 Shape 都可能改变最终结果。
最容易被忽略的正确性边界
位置编码必须跟序列分片一起走
Context Parallel 不只切 QKV。所有依赖绝对或相对 Token 位置的张量，都必须按同一 Sequence Mapping 处理，包括：
position_ids- RoPE 的
freq_cis或 Cos/Sin Cache - Attention Mask
- Packed Sequence 的边界信息
- Variable-Length Attention 的累计长度数组
- Loss Mask 与 Label Shift 后的有效 Token 范围
PyTorch 官方教程特别指出，在 Llama 类模型中若漏掉 freq_cis 的序列切分，会导致旋转位置编码计算错误。此类错误通常不会立刻 OOM 或报异常，而是表现为 Loss 偏移、长序列退化或不同 CP Size 下结果不一致。
因果 Attention 会制造 Rank 负载不均
对于 Causal Attention，注意力矩阵是下三角结构。若简单按连续 Token 区间分配，靠前的 Query 只能看到较少历史 Token，靠后的 Query 需要处理更长前缀，不同 Rank 的有效计算量并不相等。
这会形成典型的”所有 GPU 等最慢 Rank”问题。可选治理方式包括：
- 使用交错或 Striped Token Placement，让每个 Rank 同时持有早期和晚期 Token；
- 对因果 Block 做 Zig-Zag 式映射，使有效 Attention Block 更均匀；
- 使用能够跳过无效上三角 Block 的 Kernel；
- 监控每个 Rank 的有效 QK Block 数，而不是只看 Token 数。
因此，Token 数均衡不等于计算量均衡。这是 Context Parallel 上线时最重要的认知之一。
变长序列会把静态平衡重新打乱
真实训练批次通常含有不同长度样本。Sequence Packing 能减少 Padding，但也可能让不同 Pack 的有效 Block 数差异很大。若只按最大序列长度静态切分，某些 Rank 会处理大量 Padding 或被 Mask 的区域。
建议将 CP 调度与 Length Bucket、Packing Manifest 和 Variable-Length Attention 元数据联动，并记录每个 Rank 的有效 Token、有效 Attention Block 和通信字节数。
工程落地：不要从最大 CP Size 开始
第一步：建立单卡数值基线
先在单卡或 CP=1 下固定以下制品：
- 模型权重与 Tokenizer 指纹；
- Attention Backend 和精度；
- 固定输入 Token、Position ID、Mask；
- Forward 输出、Loss 和关键层梯度摘要；
- CUDA、NCCL、PyTorch、Transformer Engine 版本。
随后分别运行 CP=2、CP=4，对比完整输出和梯度。BF16/FP16 不应强求逐位一致，但必须定义可接受的绝对误差、相对误差和训练 Loss 漂移范围。
第二步：独立验证 Attention
不要直接用完整训练任务定位 CP 错误。先构造一个小型 SDPA 用例，比较单卡 Attention 与 Context Parallel Attention：
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import sdpa_kernel, SDPBackend
from torch.nn.functional import scaled_dot_product_attention
mesh = init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("cp",))
# q、k、v 的序列维位于 dim=2
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
with context_parallel(
mesh,
buffers=(q, k, v, rope_freqs),
buffer_seq_dims=(2, 2, 2, 0),
):
output = scaled_dot_product_attention(q, k, v, is_causal=True)上例仅表达配置思路。PyTorch Context Parallel API 仍处于不稳定接口阶段，实际参数、内部模块和可用 Backend 应以所使用版本的官方文档为准。
第三步：做通信方式矩阵
至少测试以下维度：
| 维度 | 测试值 |
|---|---|
| CP Size | 1、2、4、8 |
| Rotation | All-Gather、P2P Ring、All-to-All |
| 序列长度 | 短、中、目标长度、极限长度 |
| Attention 类型 | MHA、GQA、MQA |
| 拓扑 | 单机、跨机 |
| 精度 | BF16、FP16、FP8 |
| Activation Checkpointing | 开启/关闭 |
每组测试记录：Peak Memory、Step Time、Attention Time、通信时间、重叠比例、最慢 Rank 时间、NCCL 带宽和 Loss 差异。
第四步：规划混合并行拓扑
Context Parallel 通常与 Data、Tensor、Pipeline Parallel 组合。Megatron Core 的基本关系是：
world_size = data_parallel × tensor_parallel × pipeline_parallel × context_parallelCP 并不是越大越好。每增加一个维度，都会压缩其他并行组的可用规模，并改变通信边界。常见策略是：
- 先选择满足模型参数显存的 TP/PP；
- 再用 CP 解决长序列激活显存；
- 将高频 CP 通信尽量放在节点内；
- 将跨节点带宽留给频率更低或更可控的并行维度；
- 用实测决定 CP 与 Activation Checkpointing 的组合，而不是全部打开。
适用场景
Context Parallel 更适合以下工作负载：
- 长文档、代码仓库和多轮轨迹的长上下文继续预训练；
- 视频、音频、机器人轨迹等 Token 数很高的多模态训练；
- 单卡已经因激活显存 OOM，但模型参数仍能通过现有 TP/PP 放下；
- 全量 Activation Checkpointing 的计算损失过高，希望用更多 GPU 换回吞吐；
- 需要在不近似 Attention 的前提下扩展序列长度。
它不一定适合短序列训练、网络带宽有限的松散集群，或 Attention 在总耗时中占比很低的任务。此时通信启动开销可能超过显存收益。
常见误区
| 误区 | 事实 |
|---|---|
| CP 等于把输入切成几段独立训练 | 各分片仍属于同一条序列，Attention 必须跨分片交换 KV |
| 显存下降 C 倍，吞吐就一定提高 C 倍 | 激活显存可接近按 CP Size 分摊，但通信、同步等待和 Kernel 效率都会影响吞吐 |
| 所有序列张量只需本地 Slice | RoPE、Mask、Packed Sequence Metadata 等必须保持同一全局位置语义 |
| 通信带宽高就不需要看拓扑 | CP 通信频率高，跨 NUMA、跨 PCIe Switch 的路径差异会被每层重复放大 |
| 只验证短序列数值 | 很多 Position、Mask 和 Block Mapping 问题只在跨越分片边界时暴露 |
上线检查清单
正确性
- CP=1 与 CP>1 的 Forward、Loss、Gradient 已完成回放比较
- Position ID、RoPE、Mask、Label 与 Pack Boundary 使用统一 Sequence Mapping
- Causal Block 跳过逻辑经过跨分片边界测试
- MHA、GQA、MQA 的 Head/KV Head 可整除性已检查
- 变长序列和尾部不足一个 Shard 的处理规则明确
- Attention Backend 变更有独立数值门禁
性能
- 记录各 Rank Peak Memory，而不是只看 Rank 0
- 记录最慢 Rank 的 Attention 与通信耗时
- 对比 All-Gather、Ring P2P 和 All-to-All
- 验证通信是否真正与 Attention Kernel 重叠
- 验证短序列下 CP 自动关闭或降级是否更优
- 对目标长度进行至少一次完整训练 Step 压测
稳定性
- NCCL Timeout、异步错误处理与故障 Rank 日志已启用
- 保存 CP/TP/PP/DP 拓扑和版本指纹
- Checkpoint 能在目标并行拓扑下恢复
- 灰度期间保留 CP=1 或较低 CP Size 的回退路径
- 监控每 Rank 有效 Token、有效 Attention Block 和通信字节数
- 变更 CP Size 前执行自动正确性与性能回归
参考资料
- NVIDIA Megatron Core — Context Parallelism
- PyTorch Tutorials — Introduction to Context Parallel
- DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
- Ring Attention with Blockwise Transformers for Near-Infinite Context
- Striped Attention: Faster Ring Attention for Causal Transformers
- Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking