用最优传输解决 MoE 训练负载失衡：新方法 TAOT 提速 1.43 倍

大语言模型广泛采用的「混合专家（Mixture-of-Experts，MoE）」架构在提升参数规模与推理能力的同时，也带来了专家并行训练中的负载失衡难题。arXiv 最新论文《Solving MoE Load Imbalance in LLM Training via Optimal Transport》提出了一种名为 TAOT 的拓扑感知最优传输方法，在不损失均衡质量的前提下显著降低跨节点通信开销，端到端 MoE 训练速度提升 1.43 倍。

研究背景：MoE 训练中的负载失衡

MoE 模型通过路由器把每个 Token 动态分配给若干「专家」子网络，从而在保持激活参数规模可控的同时扩展总参数量。然而，路由的动态性会导致热门专家所在节点过载、空闲节点闲置，造成专家并行训练效率低下。

现有的动态副本（dynamic-replica）方案会把热门专家复制到空闲 rank 上分担计算，但这些方法只追求负载均衡数值上的最优，没有考虑多节点拓扑中专家权重迁移本身带来的通信开销。当跨节点通信代价超过均衡带来的收益时，反而会推高训练成本。

TAOT：把副本放置建模为最优传输

TAOT 的核心思路是将热门 rank 的「过载」与空闲 rank 的「富余算力」视为供需双方，用通信代价矩阵作为运输成本，构造一个带熵正则的均衡最优传输问题。

具体而言，TAOT 使用 Sinkhorn-Knopp 迭代求解该最优传输问题，得到 rank 级别的「流量提示」，再与整数副本匹配和 Token 分配结合，输出可直接执行的调度计划。在系统层面，TAOT 通过将「客座」专家权重传输与「本地」专家计算重叠执行，把通信开销隐藏到计算背后。

关键实验结果

论文报告的实验数据显示，TAOT 在多个配置下取得了全面优势：

• 端到端 MoE 训练速度提升 1.43 倍；
• 负载均衡质量与现有最优方法持平或更优；
• 加权专家通信代价在所有测试配置中最低，最高降幅达 74%。

这意味着 TAOT 在不牺牲均衡效果的前提下，大幅压缩了副本迁移带来的额外通信成本，对大规模 MoE 训练集群的效率有直接价值。

意义与适用场景

TAOT 面向的是使用专家并行训练 MoE 大模型的工程场景，例如 GPT、Mixtral、DeepSeek-V3 等采用稀疏激活架构的模型训练流程。随着主流大模型越来越多地采用 MoE 来平衡容量与算力，训练侧的通信效率将成为制约迭代速度的关键瓶颈。TAOT 通过在算法与系统两层同时优化，为这一问题提供了一个可直接落地到分布式训练框架的方案。

论文已以 arXiv:2608.03676 公开预印本形式发布，归属分布式、并行与集群计算方向，感兴趣的读者可前往 arXiv 查阅全文与方法细节。