专家并行动态负载均衡算法(数参互寻)

背景与挑战

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

解决方案

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

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

使用场景

本特性适用于以下场景：

大规模MoE模型专家并行训练，输入数据分布动态变化、专家负载不均衡显著的训练任务，并且对训练吞吐量与系统资源利用率有较高要求的分布式训练场景。

使用方法

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

--balanced-moe-experts --balanced-moe-hot-expert-num N

其中，需替换为每层期望设定的热专家数量（一个正整数）。参数说明：

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

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

使用限制

为保障专家并行动态负载均衡算法的稳定运行与性能表现，请注意以下使用限制条件：

依赖特性要求

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

分发器支持

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

并行配置限制

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

内存与模型配置

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

MoE模式限制

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

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

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

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

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

注意事项：请在使用前仔细核对上述配置要求，确保训练环境满足所有限制条件。如有特殊需求或遇到兼容性问题，请联系技术支持团队获取进一步协助。

使用效果

在专家并行训练环境中，当面临不同专家间计算与通信负载严重不均衡的挑战时，启用本特性后可实现以下优化效果：

负载均衡优化：通过数据与计算的智能互寻机制，实现专家节点的动态负载均衡。
系统吞吐量提高：基于通信量优化、计算-通信重叠及流水线掩盖等技术，提升训练迭代速度。
资源利用增强：充分利用所有专家计算资源，减少空闲等待，提升整体计算效率。