自动并行策略搜索 Expert Parallelism 估算能力支持

本文档描述自动并行策略搜索中 Expert Parallelism，EP，估算能力的设计需求。不包含具体代码实现，侧重 MoE 场景下的专家显存、token dispatch/combine、专家计算、负载不均衡建模与验收标准。

1. 背景

在 MoE 模型中，模型参数量主要来自大量专家层。Expert Parallelism，EP，通过沿专家维切分专家参数，使不同设备持有不同专家，从而降低单卡专家参数和优化器状态显存。

但 EP 也引入了 MoE 特有的 token dispatch/combine 通信，并且实际性能会受到 tokens_per_expert 负载分布影响。对于大规模 MoE 模型，自动并行策略搜索需要估算不同 ep_degree、num_experts、top_k、capacity factor 和 token 分布下的显存与性能代价。

2. 目标与非目标

2.1 目标

• 支持 ep_degree 进入自动并行策略搜索空间；
• 支持 MoE expert parameters 的显存估算；
• 支持 shared parameters 与 expert parameters 的区分建模；
• 支持 expert optimizer state 与 gradient 显存估算；
• 支持 token dispatch/combine 通信开销估算；
• 支持 expert compute cost 估算；
• 支持 tokens_per_expert 均衡和非均衡两类场景估算；
• 支持 EP 与 TP 组合下 Expert Tensor Parallel 的基础估算；
• 支持 EP 与 PP 组合时 expert layer 所在 stage 的显存和计算估算；
• 补充单元测试和 MoE toy model 集成测试。

2.2 非目标

• 本 issue 不实现 MoE router；
• 本 issue 不实现 DeepEP、HybridEP 等 runtime 优化；
• 本 issue 不实现 ExpertTensorParallel runtime；
• 本 issue 不负责 MoE 负载均衡 loss 或 expert_bias 更新逻辑；
• 本 issue 不保证精确预测训练过程中的真实 token 分布，首期支持均衡假设和用户输入分布两类估算模式；

3. 建模语义

3.1 显存估算

EP 显存估算应覆盖：

• expert weights 在 EP degree 下的切分；
• expert gradients 显存；
• expert optimizer state 显存；
• shared parameters 与 expert parameters 的区别；
• router / gate 参数显存；
• top-k routing、capacity factor 对 token buffer 的影响；
• EP 与 TP 组合时 Expert Tensor Parallel 的显存变化；
• EP 与 PP 组合时每个 stage 的 expert memory。

3.2 性能估算

EP 性能估算应覆盖：

• router / gate 计算开销；
• token dispatch all-to-all 通信开销；
• expert grouped GEMM 或专家 MLP 计算开销；
• token combine all-to-all 通信开销；
• tokens_per_expert 不均衡对 step time 的影响；
• EP 与 DP/TP/PP 组合时的通信组关系；
• 均衡 token 分布和用户指定 token 分布两类估算模式。

3.3 约束

应检查：

• num_experts 是否可被 ep_degree 整除；
• hidden size / expert hidden size 是否满足 TP 切分约束；
• top_k、capacity factor、sequence length、micro batch size 等是否形成合法 token dispatch 规模；
• dp × tp × pp × ep 等组合不超过可用设备数；
• EP 与 PP 组合时 expert layer 所在 stage 的显存和计算是否可行；
• EP 策略下单卡显存是否满足 memory limit。

4. 实现要点

• EP 估算模块应能识别 MoE 层和 dense 层；
• memory breakdown 中应区分 expert memory 和 shared memory；
• performance breakdown 中应区分 router、dispatch、expert compute 和 combine；
• token 分布可采用均衡假设，并支持用户输入分布；
• 对不支持的 EP 组合应返回 NotImplemented 或 infeasible 原因；
• 估算结果应能与 DP/TP/PP/CP 等其它维度组合。

5. 测试设计

6. 验收标准

• ep_degree 可进入自动并行策略搜索空间；
• EP expert memory 估算能力可用；
• EP dispatch/combine 通信估算能力可用；
• EP expert compute 估算能力可用；
• tokens_per_expert 均衡和非均衡两类场景可被估算；
• EP 与 TP/PP 的基础组合可被评估；
• num_experts 不满足整除约束时可被过滤；
• AP-EP-01～AP-EP-07 测试通过；
• 文档中包含 EP 建模范围、约束条件和 MoE 典型配置示例。

7. 参考

• MoE 并行训练：Expert Parallelism、Expert Tensor Parallelism；
• Megatron-LM：MoE Expert Parallelism；
• TorchTitan：ExpertParallel / ExpertTensorParallel 相关语义；
• HyperParallel：expert_parallel、MoE、自动并行策略搜索相关模块。