专家并行动态负载均衡（数参互寻）

背景与挑战

混合专家（MoE）模型通过将不同输入令牌路由至相应专家网络进行计算，以实现模型规模的扩展。然而，在实际训练过程中，输入数据分布的不确定性、路由器初始偏好以及并发请求的随机性，常导致不同专家节点间的计算与通信负载严重失衡。部分专家可能因过载成为系统瓶颈，而其他专家则处于利用率不足的状态。这种负载不均问题显著降低了训练效率与系统整体吞吐量，成为制约MoE模型规模化训练的关键挑战。

解决方案

本特性设计了一套动态自适应的专家并行负载均衡机制。其核心思想为"数据找计算"与"计算找数据"相结合，通过动态决策实现计算与参数的相互寻优，达到负载均衡。具体实现包括以下关键创新：

• 实时监控与决策：在前向传播过程中，通过路由器获取全局令牌分布信息。
• 热专家动态选择：基于全局负载信息，动态识别负载最重的专家子集作为"热专家"。
• 参数广播与共享：通过跨节点广播操作，确保所有专家并行（EP）节点均持有热专家参数的最新副本，实现参数共享。
• 令牌分发优化：重构令牌分发策略，将原本需发送至远端热专家的令牌转为本地计算，减少通信开销。
• 计算-通信重叠与流水线掩盖：在流水线并行中，通过精细调度不同Micro-batch的前向与反向计算阶段，实现专家并行通信的完全掩盖。

该方案通过"数据找计算"与"计算找数据"相结合的动态决策机制，实现计算与参数的智能互寻，最终达到负载均衡，提升系统吞吐量和资源利用率。

使用场景

本特性适用于以下场景，强烈建议同时满足以下两个条件时启用，否则可能无法获得性能收益甚至出现负收益：

1. 大规模专家并行（EP）训练：

   • EP原本配置较大（例如≥32）：本特性的收益与专家并行度的规模正相关。当EP值较小时，负载不均衡的程度相对有限，而启用本特性引入的参数广播等额外开销可能超过其带来的收益。在EP值较大时，负载不均衡问题通常更突出，通过本特性优化后获得的性能提升也更显著。

2. 负载显著不均衡：

   • 在实际训练中，输入数据分布动态变化导致部分专家接收的令牌数量远超其他专家，造成明显的计算与通信热点。这种场景下，通过动态负载均衡重新分配计算任务能带来最大收益。

不推荐使用或收益有限场景：

• EP配置较小（例如≤8）且负载相对均匀：额外通信开销可能抵消负载均衡的收益。
• 无法满足下述"使用限制"中任一条款的配置。
• 需要严格确定性训练的场景：本特性涉及动态决策和异步通信，不保证二进制对齐和完全确定性。
• DSKv3减层4机Atlas A3场景loss误差<1%，小尺寸模型场景累加序引起的误差可能会被放大，不推荐使用。

使用方法

用户可在训练脚本中通过添加以下参数，开启专家并行动态负载均衡算法：

--balanced-moe-experts --balanced-moe-hot-expert-num N --trans-hot-expert-group-num M

参数说明：

• --balanced-moe-experts：总开关，用于启用动态负载均衡算法。

• --balanced-moe-hot-expert-num <N>：指定每层动态维护的热专家数量 N（一个正整数）。系统将根据此参数及实时负载状况，选择负载最大的 N 个专家进行参数广播与负载均衡调度。

  • 默认值： 3
  • 重要建议与限制：
    1. N 必须小于等于每个EP Rank上的本地专家数。本地专家数计算公式为：num_local_experts = num_experts / expert_model_parallel_size。例如，总专家数为64，EP度为8，则每个Rank有8个本地专家，此时 --balanced-moe-hot-expert-num 最大只能设置为8。
    2. 设置过大的 N（超过本地专家数）不仅不会带来额外收益，反而会增加不必要的参数广播开销和内存占用，可能导致性能下降。建议根据实际负载不均程度，将 N 设置为（通常3-8之间），以在开销与收益间取得平衡。

• --trans-hot-expert-group-num <M>：指定传输热专家参数时的分组数量 M（一个正整数）。用于控制热专家参数广播的并发度，有助于减少通信时间。

  • 默认值： 3
  • 重要限制： M 必须满足：1 ≤ M ≤ N（即--balanced-moe-hot-expert-num的值）。

使用限制

为保障专家并行动态负载均衡算法的稳定运行与性能表现，请严格遵守以下使用限制与配置要求：

1. 依赖特性要求

• 必须同时启用 --moe-fb-overlap 参数。
• 必须启用 --moe-grouped-gemm 以支持GroupedMatmul运算。

2. 分发器支持

• 当前仅支持 --moe-token-dispatcher-type=alltoall 分发器类型。
• 暂不支持 allgather 及 alltoall_seq 类型的分发器。

3. 并行配置限制

• 需设置 --expert-tensor-parallel-size=1，暂不支持专家张量并行。
• 需关闭 --overlap-grad-reduce，暂不支持异步数据并行通信掩盖。

4. 内存与模型配置

• 仅支持 --moe-zero-memory=level0 内存优化级别。
• 不支持 moe-zero-memory-num-layers 配置参数。
• 仅支持Mcore架构模型，需确保关闭 --use_legacy_models。

5. MoE模式限制

• 仅支持Dropless模式，暂不支持Megatron MoE的Token Drop & Pad功能。
• 不建议同时启用 --swap-attention，可能导致性能劣化。

6. 流水线并行约束

在使用虚拟流水线并行（VPP）时，需满足以下条件：

• 全局批次大小（GBS）需满足：GBS > 1 × 数据并行度（DP） × 流水线并行度（PP） × 微批次大小（MBS）。
• 若使用noop layers，必须将其置于模型尾部的最后一个VPP阶段。

7. 冲突特性

以下特性与专家并行动态负载均衡算法存在冲突，不可同时启用：

• --moe-alltoall-overlap-comm
• --moe-hierarchical-alltoallv
• --recompute-in-advance
• --recompute-in-bubble

8. 参数配置约束

• --balanced-moe-hot-expert-num 必须小于等于 每个EP Rank上的本地专家数（num_experts / expert_model_parallel_size）。
• --trans-hot-expert-group-num 必须满足：1 ≤ M ≤ N（其中N为--balanced-moe-hot-expert-num的值）。

9. 非确定性行为说明

重要说明：本特性不是二进制对齐的

1. 计算路径的根本性改变：

   • 正向传播路径改变：开启本特性后，原本需要发送到远端热专家的令牌改为在本地计算，这改变了整个计算图的执行路径。
   • 梯度累加顺序改变：热专家的梯度会在多个EP Rank上并行计算和异步归约，改变了梯度累加的顺序。

2. 浮点数运算的非交换性影响：

   • 由于浮点数加法不具备结合律，改变令牌在不同设备间的分配顺序会导致正向传播和反向传播的中间结果产生微小差异。
   • 即使数学上等价的运算 (A+B)+C 和 A+(B+C)，在浮点数表示下可能产生不同的结果。

3. 非二进制对齐的具体表现：

   • 同一训练多次运行不一致：即使设置确定性计算标志 --npu-deterministic，运行相同的训练两次也可能不会产生比特级完全相同的结果。

4. 适用场景限制：

   • 绝对不适用：需要严格比特级可重现性的场景。

使用效果

在满足推荐使用场景（EP规模大且负载不均）并正确配置的专家并行训练环境中，启用本特性后可实现以下优化效果：

• 负载均衡优化：通过"数参互寻"机制，动态地将计算负载从过载的"热专家"迁移至空闲专家，使各专家节点的计算量趋于平衡，缓解系统瓶颈。
• 系统吞吐量有效提高：通过将原本需要跨设备通信的远端热专家令牌转为本地计算，显著降低了All-to-All通信的数据量。同时，利用精细的计算-通信重叠及流水线调度，有效掩盖了参数广播等新增通信的延迟，从而提升训练迭代速度。
• 硬件资源利用率提升：减少了因负载不均导致的部分设备过载而其他设备空闲等待的时间，使所有计算单元得到更充分的利用，提升整体集群效率。

典型性能收益：在符合推荐场景（如EP≥32、负载不均显著）且配置合理的条件下，本特性可带来显著的端到端训练吞吐量提升。反之，若不满足推荐场景或配置不当，则可能无法获得收益甚至导致性能下降。