arXiv:2605.08639cs.LG2026-05被引 2

利用RL的路由回放机制,实现MoE训练中微批次级负载均衡。

ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning

论文配图:ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning
图 1 · 摘自论文原文
  • 基于滚动-训练流程中的路由回放信息,预知专家分配。
  • 在跨批次和微批次层面分别进行专家重排序与动态复制。
  • 适用于大规模MoE语言模型的强化学习训练场景。

负载不均是混合专家(MoE)训练的长期挑战,在大语言模型的强化学习(RL)训练中尤为严重,因热点专家在微批次间频繁切换。现有系统依赖历史负载预测未来需求,难以应对剧烈波动。我们提出ReLibra,利用强化学习中滚动与训练共享相同令牌和参数的特性,提前获知每个令牌的专家路由决策。基于此,ReLibra在跨批次和微批次两个时间尺度部署负载均衡机制:跨批次通过专家重排序实现节点间负载平衡;微批次内通过动态专家复制吸收局部波动。在多种MoE LLM及RL工作负载上的实验表明,ReLibra相比Megatron-LM提升训练吞吐最高达1.6倍,相比EPLB最高提升1.2倍,即使在给定理想负载的情况下亦然。其性能仅比理想均衡基线低6%-10%。

原文摘要 · Abstract (English)

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.

MoE强化学习负载均衡大模型训练

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