arXiv:2502.15451cs.LGcs.CL2025-02被引 1

用整数规划解决专家负载不均问题,训练更高效且全程平衡。

Binary-Integer-Programming Based Algorithm for Expert Load Balancing in Mixture-of-Experts Models

  • 基于二元整数规划动态调整专家选择顺序,实现负载均衡。
  • 在0.3B和1.1B模型上,困惑度最低,预训练时间减少至少13%。
  • 首次实现从头到尾每层每个专家负载始终均衡,适合大规模MoE训练。

预训练混合专家(MoE)模型时,专家负载不均会导致路由坍塌或计算开销增加。现有方法如带损失控制和无损失控制法,在训练初期负载不均程度仍很高且下降缓慢。本文提出基于二元整数规划(BIP)的负载均衡算法BIP-Based Balancing,为每层MoE维护一个额外向量q,通过求解小规模整数规划快速调整s的top-K顺序。我们在两个MoE语言模型上实现:16专家(0.3B)和64专家(1.1B)。实验表明,相比损失控制法与无损失法,本算法训练出的模型困惑度最低,且预训练时间至少节省13%。据我们所知,这是首个在预训练全过程(从第一步到最后一步)中,保持每层每个专家负载均衡的路由算法,同时模型性能优异。代码已开源:https://github.com/sunyuanLLM/bip_routing_algorithm。

原文摘要 · Abstract (English)

For pre-training of MoE (Mixture-of-Experts) models, one of the main issues is unbalanced expert loads, which may cause routing collapse or increased computational overhead. Existing methods contain the Loss-Controlled method and the Loss-Free method, where both the unbalanced degrees at first several training steps are still high and decrease slowly. In this work, we propose BIP-Based Balancing, an expert load balancing algorithm based on binary integer programming (BIP). The algorithm maintains an additional vector q on each MoE layer that can help change the top-K order of s by solving a binary integer programming with very small time costs. We implement the algorithm on two MoE language models: 16-expert (0.3B) and 64-expert (1.1B). The experimental results show that on both models comparing with the Loss-Controlled method and the Loss-Free method, our algorithm trains models with the lowest perplexities, while saves at least 13% of pre-training time compared with the Loss-Controlled method. Within our current knowledge, this is the first routing algorithm that achieves maintaining load balance status on every expert in every MoE layer from the first step to the last step during the whole pre-training process, while the trained MoE models also perform well. The code material of this work is available at https://github.com/sunyuanLLM/bip_routing_algorithm.

MoE负载均衡整数规划预训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。