ReLibra：基于路由回放的MoE训练负载均衡
ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning
AI总结：
ReLibra通过利用强化学习的回放训练流程中的路由回放机制，在微批次粒度上实现细粒度负载均衡，提升MoE训练效率。
摘要：
负载不均衡是混合专家（MoE）训练中的长期挑战，尤其在大规模语言模型的强化学习（RL）中更为严重，因为热点专家在微批次间频繁切换。现有MoE训练系统依赖历史负载预测未来专家需求，但在剧烈波动时效果不佳。我们提出ReLibra，一种利用强化学习回放训练流程中独特机会的MoE RL训练系统，通过在微批次粒度上实现细粒度负载均衡。由于回放和训练处理相同的token并使用相同的MoE参数，token到专家的路由决策在训练开始前已知。利用这一信息，ReLibra在不同批次时间尺度上部署两个MoE负载均衡机制，匹配其通信模式以适应分层网络带宽。在跨批次时间尺度上，ReLibra执行专家重新排序以在节点层面实现跨节点平衡；在内批次时间尺度上，它在节点内动态执行专家复制以吸收微批次层面的负载波动。在多种MoE LLM和RL工作负载上的实验表明，ReLibra在Megatron-LM上提升了1.6倍的训练吞吐量，在EPLB上提升了1.2倍，即使EPLB被赋予了理想负载。此外，ReLibra的吞吐量仍接近理想均衡基线的6%-10%。
英文摘要：
Load imbalance is a long-standing challenge in Mixture-of-Experts (MoE) training and is exacerbated in reinforcement learning (RL) for LLMs, where hot experts can shift frequently across micro-batches. Existing MoE training systems rely on historical loads to predict future expert demand, making them less effective under sharp fluctuations. We propose ReLibra, an MoE RL training system that exploits a unique opportunity in RL's rollout-training workflow, routing replay, to enable fine-grained load balancing at micro-batch granularity. Because rollout and training process the same tokens with the same MoE parameters, the token-to-expert routing decisions are known before training starts. Leveraging this information, ReLibra places two MoE load-balancing mechanisms at inter- and intra-batch timescales, matching their communication patterns to hierarchical network bandwidths. At the inter-batch timescale, ReLibra performs expert reordering to redistribute experts for batch-level cross-node balancing; at the intra-batch timescale, it dynamically performs expert replication within a node to absorb micro-batch-level load fluctuations. Experiments on diverse MoE LLMs and RL workloads show that ReLibra improves training throughput by up to 1.6$\times$ over Megatron-LM and by up to 1.2$\times$ over EPLB, even when EPLB is given oracle loads. Moreover, ReLibra remains within 6%-10% of the throughput of an idealized balanced baseline.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
