基于二进制-整数规划的专家负载均衡混合专家模型算法
机器学习
2025-03-21 v2 计算与语言
摘要
对于混合专家 (MoE) 模型的预训练，主要问题之一是专家负载不均，这可能导致路由崩溃或增加计算开销。现有方法包括损失控制方法和损失自由方法，其中在前几个训练步骤中仍存在较高的不平衡度且下降速度较慢。本文提出了基于二进制整数规划 (BIP) 的负载均衡方法。该算法在每个 MoE 层维护一个额外向量 ，通过解决时间成本极小的二进制整数规划问题来帮助改变 的 Top-K 排序。我们在两个 MoE 语言模型上实现了该算法：16 个专家 (0.3B) 和 64 个专家 (1.1B)。实验结果表明，在与损失控制方法和损失自由方法进行比较时，我们的方法在两个模型上训练的模型具有最低的困惑度，同时相对于损失控制方法节省至少 13% 的预训练时间。就目前的知识而言，这是第一个在整个预训练过程中从第一步到最后一步维持每个 MoE 层中每个专家负载平衡状态的路由算法，而且训练得到的 MoE 模型也表现良好。该工作的代码材料可在 https://github.com/sunyuanLLM/bip_routing_algorithm 获取。
引用
@article{arxiv.2502.15451,
title = {Binary-Integer-Programming Based Algorithm for Expert Load Balancing in Mixture-of-Experts Models},
author = {Yuan Sun},
journal= {arXiv preprint arXiv:2502.15451},
year = {2025}
}