[Chamberlain0w0]

概述

本 PR 为 GPT-2、LLaMA 和 Qwen3 引入 Context Parallel（CP），将同一样本的序列分布到多个 rank，降低单卡激活与 attention 中间结果的显存需求，并与现有 TP/SP、PP/VPP 和 ZeRO 训练流程衔接。

主要改动

支持四种 CP 通信模式及其前向、反向传播：

模式	实现方式
all_gather	收集全局 KV，以本地 Q 计算 attention；反向通过 ReduceScatter 汇总 KV 梯度
p2p	沿 ring 交换 KV，使用单遍 online softmax 合并分块结果；反向重算并将 KV 梯度返回原 owner
a2a	通过 AllToAll 交换 sequence/head 维度，在完整序列上计算 attention，再恢复输出布局
a2a+p2p	在 CP group 内划分 A2A/P2P 子组，先执行 A2A 布局转换，再执行 ring attention

• 提供默认的 contiguous 布局和可选的 zigzag 首尾配对布局，后者用于均衡 causal attention 在不同 rank 上的计算量。
• 扩展并行拓扑和通信组，使用 DP×CP 同步复制参数的梯度；dataloader 仍按 pure DP 划分样本。
• 在模型入口统一接入 CP 参数，适配输入与 label 切片、位置编码、causal mask、梯度累积、loss 归约和 pipeline 边界形状。
• 扩展 checkpoint 的 CP 元数据与保存 rank 选择，复用现有保存恢复机制。
• 为 ring 前反向增加双计算流调度，通过 event 和通信完成依赖管理跨流执行及 KV 生命周期。

数值精度

• 修正 CUDA BF16 Softmax 的提前舍入，使概率在归一化后转换为 BF16。
• 调整 CUDA GEMM 的中间归约及 workspace 策略，减少矩阵切分改变计算路径时产生的精度差异；该策略同时影响 non-CP。
• 对齐 ring backward 的 autocast 输入和上游梯度 dtype。
• 保留单遍 ring 和按输入 dtype 通信，并通过 FIXME 说明 BF16 TP/SP 求和顺序、块内概率/PV 舍入及反向概率重算可能带来的差异。

验证

• 算子测试使用独立 FP64 reference，覆盖四种通信模式、两种布局、MHA/GQA、CP2/4/8、2D hierarchy，以及非连续 CP ranks 的输出和 Q/K/V 梯度。
• 完成 GPT-2、LLaMA、Qwen3 的多步 loss 对照，以及部分 TP/SP/PP/VPP/ZeRO、重计算和 checkpoint 保存恢复组合验证。
• 当前构建的 15 项 Softmax/Matmul 测试、19 项 CP/布局/拓扑测试通过。
• 全量 clang-format、Black、clang-tidy、IWYU 检查通过。

实现边界

当前支持固定长度、等分的 causal self-attention，尚未接入 FlashAttention/fused CP backend，也未覆盖 packed/变长序列和全部并行组合。

FP32 在已验证配置下满足数值容差；BF16 不保证跨并行配置逐位一致，在部分配置下可能会存在累加顺序不一致导致误差较大的情况。

[chen2021673]

这里影响默认精度吗，pytorch 脚本是否需要同步修改 torch.backends.cuda.matmul.allow_<dtype>_reduced_precision_reduction

[Chamberlain0w0]

会影响，需要设置：

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False

[Chamberlain0w0]

11.13 更新：不在此 PR 中做任何修改非 CP 路径的数值精度策略，已经恢复

[chen2021673]

这个函数最好把 RepeatKVHeads 放在内部执行，因为之后这部分拷贝应该可以通过改 kernel 省下来（可以加个 TODO），移到内部的话可以只改 kernel 调用不改接口。参数可以和 megatron 对齐，传 num_gqa_groups 来推 H_q/H_kv。
https://github.com/NVIDIA/TransformerEngine/blob/c1c2d4d2fa008386f247c2fec56a998e27134e0a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py#L2781
另外，q@k^T × 1/√D → MaskedFill → Softmax → @v 这套逻辑 causal_self_attention 里也有，这部分能不能直接抽出来改成 ScaledDotProductAttention 共用

[Chamberlain0w0]

是的，这块 FlashAttn 那个 PR 里是提了一个 ScaledDotProductAttention 的 function 的，但是这个 PR 里主要是想不太动原先的小算子拼接的逻辑。我这块留个 FIXME 吧，我觉得可以放在之后哪个 PR 大改 attn 的时候再顺手改一下？

[Chamberlain0w0]

我直接在这个 PR 里改了吧。FlashAttn 那个可能还要再往后放，我看了下改动应该不大，主要还是整块整块的结构上的变动

[chen2021673]

这里叫 num_gqa_groups， Megatron 同名参数语义是 H_kv，但传的是 n_rep = H_q/H_kv，建议统一一下

[Chamberlain0w0]

改了

[Chamberlain0w0]