[Chamberlain0w0]

目标

为启用 sequence parallel 的 TP Linear 增加通信与计算重叠，支持 backward bulk overlap，以及 forward/backward 的分块 AG/RS 流水调度。GPT-2、LLaMA3、Qwen3 均提供对应命令行选项，总开关默认关闭。

主要修改

• 新增 ModelParallelConfig，统一保存模型并行拓扑与 overlap 配置，由 GlobalEnv 持有并提供只读访问，同时保留原有初始化接口。
• 使用一个 LinearWithGradAccumulationAndAsyncCommunication Function 承接 Column/Row 的 SP 计算和通信，同步与重叠路径共用实现。
• 将 Linear 的原始 dispatcher 调用整理为 autograd::linear 接口，供普通 Linear 和 TP Function 复用，并为 TP Function 注册 autocast 策略。
• 使用现有 BatchSendRecv 实现 ring AG/GEMM 流水；新增异步 ProcessGroup::Reduce，通过按目标 rank 分块归约实现 GEMM/RS 流水。
• 为 QKV/FC1 设置层标识，以支持按投影禁用流水；LoRA 包装和合并时保留该标识。
• 正确处理普通训练、NoGrad、reentrant 和 non-reentrant checkpoint 的上下文保存与临时张量释放。

配置选项

所有重叠选项在 sequence_parallel=true 且 TP size > 1 时生效。tp_comm_overlap=false 时不启用重叠。

选项	默认值	控制行为
tp_comm_overlap	false	重叠总开关
tp_comm_bulk_wgrad	true	Column backward：激活 AG 与 dgrad GEMM 重叠
tp_comm_bulk_dgrad	true	Column backward：dgrad RS 与 wgrad GEMM 重叠
tp_comm_overlap_ag	true	Column forward AG/GEMM、Row backward AG/dgrad 分块流水
tp_comm_overlap_rs	true	Row forward GEMM/RS 分块流水
tp_comm_overlap_rs_dgrad	false	Column backward dgrad/RS 分块流水
tp_comm_overlap_disable_qkv	false	禁用 QKV 的 forward AG 和 split dgrad RS 流水
tp_comm_overlap_disable_fc1	false	禁用 FC1 的上述两条流水

QKV/FC1 的 disable 选项遵循 Megatron TE 包装层的控制语义，不关闭 bulk overlap。

SP Column 为 wgrad 保存本地 sequence shard，在 backward 重新执行 AG；这一激活保存策略在同步和重叠模式下保持一致。Row 保存覆盖完整 sequence 的本地输入特征 shard，bias 在 RS 后处理。

验证

• GPT-2、LLaMA3、Qwen3、Mixtral 构建通过。
• 最新 master 上的现有回归：227 项通过、2 项设备条件跳过，覆盖 autograd、checkpoint 重分片、StreamGuard、LoRA、Transformer/Qwen 和并行拓扑。
• 本地专项数值验证覆盖 FP32/BF16、TP=1/2/8、所有子开关组合，以及 NoGrad、两种 checkpoint、冻结梯度和 bias 场景。
• 补充验证共享权重、残差连接、微批次梯度累积、单机多进程，以及 StreamGuard 非默认流执行。
• 全量 clang-format、Black、clang-tidy 和 IWYU 本地检查通过；clang-tidy/IWYU 使用与 CI 一致的 CPU-only 配置。

实现范围

当前基于 InfiniTrain 现有通信流、event 和 GEMM 接口实现 bulk 与分块流水调度，不包含 TE userbuffer、通信/GEMM 融合 kernel 或 workspace/event pool。

分块会增加 GEMM、通信调用和临时 buffer 开销，实际收益依赖模型、形状和硬件；全部开启不一定最快，需要针对训练配置调优。

[chen2021673]

这里就算关了 tp_comm_overlap 也会进入新的 SP Linear 实现，影响了默认行为。在新的 linear backward 里，wgrad 需要重新 AllGather input，而不是保存 forward 中 Gather 后的完整 input，这会增加通信成本。虽然理论上可以并行起来，但建议还是实测一下，贴下结果。

[chen2021673]

哦不对，总开关关闭时，所有子开关不生效，tp_comm_overlap 关了的话 tp_comm_bulk_wgrad 也是关的，不会重叠，那不开会对 sp 的性能有影响吗

[Chamberlain0w0]

差别其实在于之前实现的 SP，attn/mlp 部分的激活值会保存全部 seq，但实际上开了 SP 应该在这块只保存 local seq 就够了，这就在反向中会引入额外 allgather 通信。Megatron 里面会默认自动把这个通信跟计算隐藏起来。

[Chamberlain0w0]

Megatron 里面的 dgrad 和 wgrad 的 bulk_overlap 的两个选项（tp_comm_bulk_dgrad 、tp_comm_bulk_wgrad）其实不是控制是否在 backward 做 overlap（因为默认路径就是 overlap 的了），而是指是否调用 te 的后端的优化来实现这个重叠。但以我们目前框架的状况，我还是把这两个选项解读为是否显式开启/关闭 bulk 重叠，我这里加个 FIXME 吧。

[chen2021673]

这里最好不要硬编码

[Chamberlain0w0]

改了一下，把判断移到创建层的位置，ColumnParallelLinear 构造函数末尾参数改为 bool disable_tp_comm_pipeline = false，替换当前字符串参数，只接收最终的禁用标志。Attention 创建 QKV 时传入 config.tp_comm_overlap_disable_qkv；MLP 创建 FC1 时传入 config.tp_comm_overlap_disable_fc1。这里的 config 指全局 ModelParallelConfig。然后 ColumnParallelLinear 保存这个 bool，在 Forward 中直接透传给 Function 的 disable_pipeline，删除字符串判断。

[chen2021673]

确实，tp_comm_buffer_name 在megatron里还要给TE用，在这里不用，所以可以不传具体参数。麻烦修改一下 commit 313fe05 的拼写错误，并补充性能、精度测试结果贴图。