MoE 训练中的长尾问题

MoE 已成为大模型的主流架构。每个 token 只激活少数专家，模型规模可以继续扩大，而单步计算量不必同步增长。但专家并行（EP）存在一个现实问题：每张卡需要处理多少 token，完全由路由结果动态决定。少数专家一旦成为热点，token 便会集中到承载这些专家的 rank，整体训练耗时受这几张卡限制，其余卡完成计算后只能空闲等待。

现有方案忽略的因素

业界很早就提出了应对方案：某个专家的计算量过高时，临时复制一份它的权重，将副本放置到空闲卡上，由该卡分担部分计算。 问题在于，这些方案几乎只优化负载是否均衡，并默认副本可以放置到任意空闲卡。但真实硬件环境并非如此：

  • 同一台机器内的 GPU 通过 NVLink 高速通道互联，数据传输速度快、开销低；
  • 跨机器的 GPU 需要通过网络（InfiniBand）互联，传输相同规模的数据时速度更慢、开销更高。

这会带来一个容易被忽略的后果：两个方案即使负载均衡程度相同，只要一个将权重保留在本机、另一个将权重传输到其他机器，实际通信开销就可能存在显著差异。

TAOT 的思路：同时考虑传输代价

百度 LoongForge 团队认为，副本的放置位置不能只取决于空闲容量，还需要评估传输代价。 TAOT 正是从这里入手：它首次将**“削峰收益”和“跨节点传输权重的代价”同时纳入优化目标**。选择副本放置位置时，既要实现负载均衡，又要尽量降低跨机器传输的通信开销。 它借鉴了数学中的“最优传输”（Optimal Transport）思想，将通信代价纳入调度决策：

  • 优先放置在本机：本机存在空闲容量时，优先在本机放置副本，减少跨节点传输；
  • 必要时跨机器放置：本机容量不足时，仍允许跨机器放置，但将其视为代价更高的选项，而不是完全禁止。

现有方案基本落在两个极端。一类完全不考虑拓扑，只要有空闲卡就放置副本，副本容易被放置到其他机器上；负载虽然得到均衡，但跨机通信开销显著增加。另一类则使用固定拓扑图限制权重传输路径。规模扩大后，热点被拆散，近处容量先被用尽，远处仍有空闲容量的卡却可能因拓扑约束无法使用，可调度空间反而受限。 TAOT 采取折中策略：不禁止跨机传输，也不固定路径，而是为跨机传输赋予更高代价，使其在优化时自然排在后面。因此，规模越大、空闲卡越多，TAOT 的调度空间越大。系统整体架构如图所示：

TAOT 系统整体架构图

TAOT 的整体流程如下：路由完成后，先收集每张卡的负载情况，交给三阶段规划器确定副本方案；执行时，再将副本权重传输到目标卡，并使传输与 GPU 计算并行执行。 “均衡”和“通信代价”两个目标联合优化时，问题规模会随 EP 度数迅速膨胀。TAOT 将其拆分为三个阶段，由粗到细，分别解决一个问题：

  • Phase 1（rank 级）：负载大致流向何处？将过载量作为供给、空闲容量作为需求，结合通信代价矩阵求解最优传输，得到全局流量建议表。该表只是软提示，为后续阶段提供全局参照，避免局部决策。
  • Phase 2（专家级）：哪个专家的副本进入哪张卡？将连续流量转换为 0/1 决策：为每个“（空闲卡，热点专家）”配对计算分数，分数由削峰收益、放置距离和 Phase 1 流量提示加权得到，选择得分最高的组合。
  • Phase 3（token 级）：这些 token 分别由哪张卡发送？同一个专家的 token 分布在多张卡上，通过多轮竞价确定每张源卡的发送量。中标者价格上涨，从而在后续轮次中降低优先级，将就近原则融入分配过程。

算法之外的工程实现

要将方法集成到训练流程中，除了算法本身，还需解决两个工程问题。 第一，降低调度开销。这套规划需要在训练过程中逐个 microbatch 实时执行。若规划本身耗时较高，通信节省的时间就会被调度开销抵消。为此，团队在算子层面进行了优化，将原本需要执行数百次的零散计算压缩为几次，最终将规划开销控制在前向计算时间的 1% 以内。 第二，将权重传输与计算重叠。副本权重的分发在通信流上执行，与同一张卡上的 home 专家 GEMM 并行；反向阶段的 guest 梯度通过一次 All-to-All 回传并累加。传输时间基本被计算掩盖。

实验效果

  • 更快：在 Qwen3-30B-A3B 上，单次迭代的前向与反向计算时间从 155.4ms 降至 108.8ms，端到端加速 1.43 倍；
  • 更省：在负载均衡程度相当的前提下，TAOT 的通信代价在对比的 SOTA 方法中最低，最多比对手低 74%；
  • 不牺牲精度：该加速来自更合理的调度，而非降低数值精度，loss 曲线与标准方案基本重合；
  • 规模效应：并行规模越大、初始负载越不均衡，TAOT 带来的加速越显著。EP4 到 EP16 的加速最高为 1.79×；初始不均衡度从 30% 升至 90% 时，加速从 1.21× 提升至 1.75×。在负载更不均衡的场景中，收益更大。

结语：通信代价感知的调度器

动态副本是 MoE 负载均衡中的成熟技术路线，Echo、LPLB 和 LLEP 均属于该路线。TAOT 的核心差异在于优化目标：将副本放置位置的通信代价与削峰收益纳入同一个目标函数进行权衡，而不是只追求负载均衡。实现上使用连续代价矩阵和软拓扑偏好，将“优先节点内、必要时再跨节点”直接纳入目标，无需特殊硬件、无需限制在单节点，也无需预定义拓扑图。 TAOT 可用于缓解大规模 MoE 训练中的热点专家长尾问题。TAOT 已作为 MoE 专家机制扩展，集成到百度百舸开源的全模态训练框架 LoongForge 中。