自动并行策略搜索 Context Parallelism 建模能力支持

本文档描述自动并行策略搜索中 Context Parallelism，CP，建模能力的设计需求。不包含具体代码实现，侧重长序列场景下的显存、通信、性能建模与验收标准。

1. 背景

在长序列大模型训练中，attention 相关 activation 显存会随 sequence length 增长而快速上升，成为训练长上下文模型的重要瓶颈。

Context Parallelism，CP，通过沿 context / sequence 维切分长序列，使每个 rank 持有或计算部分上下文相关状态，从而降低单卡 attention activation 显存压力。但 CP 也会引入额外的 attention 通信，例如 KV exchange 或 ring attention 类通信。

自动并行策略搜索需要将 CP 纳入搜索空间，并估算其对长序列模型显存和性能的影响。

2. 目标与非目标

2.1 目标

• 支持 cp_degree 进入自动并行策略搜索空间；
• 支持 CP 对 attention activation 显存的影响估算；
• 支持 CP 引入的通信开销估算；
• 支持 CP 与 TP/PP 的基础组合建模；
• 支持 sequence length 与 cp_degree 的整除约束校验；
• 支持长序列场景下 CP 策略的可行性评估；
• 补充单元测试和估算趋势测试。

2.2 非目标

• 本 issue 不实现新的 CP runtime；
• 本 issue 不实现具体 ring attention 算子；
• 本 issue 不负责 PP 负载均衡建模；
• 本 issue 不保证覆盖所有 attention 变体。

3. 建模语义

3.1 内存估算

CP 内存估算应覆盖：

• CP degree 对 attention activation 显存的影响；
• 长序列下 KV / attention score / softmax 等中间激活的切分效果；
• CP 与 TP 组合时的 activation layout 变化；
• CP 引入的通信 buffer 或临时 tensor 显存。

3.2 性能估算

CP 性能估算应覆盖：

• CP 引入的 attention 通信开销；
• ring attention 或 KV exchange 类通信模式的近似代价；
• 不同 sequence length 和 cp_degree 下的收益与开销趋势；
• CP 与 TP/PP 组合时的通信边界变化；
• CP 对端到端 step time 的影响。

3.3 约束

应检查：

• sequence length 是否满足 cp_degree 的切分约束，例如整除或可 padding 对齐；
• cp_degree 是否超过可用设备数；
• CP 与 TP/PP 的组合是否合法；
• CP 是否仅在长序列或指定模型结构下启用；
• CP 策略下单卡显存是否满足 memory limit。

4. 实现要点

• CP 估算模块应作为独立 estimator 接入自动并行策略搜索器；
• 应输出 CP 对 memory breakdown 的影响；
• 应输出 CP 引入的 communication cost；
• 对不支持的 CP 组合应返回 NotImplemented 或 infeasible 原因；
• 应支持通过配置开启或关闭 CP 搜索；
• 估算结果应能与 DP/TP/PP 的估算结果组合。

5. 测试设计

6. 验收标准

• cp_degree 可进入自动并行策略搜索空间；
• CP attention activation memory 估算能力可用；
• CP 通信开销估算能力可用；
• CP 与 TP/PP 的基础组合可被评估；
• sequence length 不满足整除约束时可被过滤；
• AP-CP-01～AP-CP-06 测试通过；
• 文档中包含 CP 建模范围、约束条件和典型长序列配置示例。

7. 参考

• 长序列训练：Context Parallelism / Ring Attention 类方法；
• Megatron-LM：Context Parallel 相关实践；
• HyperParallel：自动并行策略搜索与内存估算相关模块。