ReLibra：基于路由-回放驱动的混合专家（MoE）强化学习训练负载均衡
机器学习
2026-05-12 v1
摘要
负载不平衡是混合专家（Mixture-of-Experts, MoE）训练中长期存在的挑战，在大语言模型（LLM）的强化学习（RL）场景中尤为突出，因为热点专家可能在不同微批次间频繁迁移。现有MoE训练系统依赖历史负载预测未来专家需求，在面对剧烈波动时效果受限。我们提出ReLibra，一种利用RL rollout-training工作流程中独特机会——路由回放（routing replay）——实现微批次粒度细粒度负载均衡的MoE RL训练系统。由于rollout与训练处理相同 token且使用相同的MoE参数，训练开始前即可知晓token到专家的路由决策。基于此信息，ReLibra在跨批次和跨批次内两个时刻放置两种负载均衡机制，匹配其通信模式与层次化网络带宽。跨批次时，ReLibra进行专家重排以实现跨节点批级负载均衡；跨批次内时，动态进行专家复制以吸收微批次级负载波动。在多样化MoE LLM和RL工作负载上的实验表明，ReLibra相较于Megatron-LM提升至1.6倍，相较于EPLB提升至1.2倍，即便EPLB获准确负载预测也未能超越。此外，ReLibra仍在理想平衡基线的6%-10%范围内。
引用
@article{arxiv.2605.08639,
title = {ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning},
author = {Chao Jin and Xinming Wei and Yinmin Zhong and Chengxu Yang and Bingyang Wu and Ruidong Zhu and Zili Zhang and Yuliang Liu and Xin Jin},
journal= {arXiv preprint arXiv:2605.08639},
year = {2026}
}