Summary

This paper introduces ForeMoE, a novel system for efficient expert load balancing in MoE (Mixture-of-Experts) models during Reinforcement Learning (RL) post-training. The authors identify that existing load-balancing techniques are inadequate for the unique characteristics of RL post-training and propose a system that leverages the foreseeability of routing decisions to achieve significant performance gains.

Summary (Overview)

• Problem: MoE models suffer from load imbalance, causing significant efficiency bottlenecks during RL post-training (e.g., PPO, GRPO). Existing systems designed for pre-training fail to address the specific challenges of this phase.
• Key Insight: RL post-training allows access to routing decisions before execution in the recompute and policy update stages Mendacious? No, the summary says "This exposes...routing information in advance". This foresight enables proactive, rather than reactive, load balancing.
• Contributions:
  • Introduces a micro-step-level load-balancing framework, contrasting with the step-level granularity used in pre-training systems.
  • Proposes a hierarchical planner that decomposes the NP-hard load-balancing problem into tractable sub-problems (base placement, expert replication/relocation).
  • Develops a transfer engine that supports both CPU-assisted and GPU-direct paths, enabling low-overhead expert migration and overlapping with computation.

Theoretical Foundation

The paper identifies that while step-level expert load is stable but skewed, micro-step-level load is highly variable and skewed. This motivates the need for a finer granularity of load balancing. The inherent challenges are:

1. Stability: Step-level routing distribution changes little across training steps.
2. Variability: Micro-step-level expert loads fluctuate substantially due to small sample sizes.
3. Complexity: The joint problem of expert placement and token routing is NP-hard.

Key Results

• Hardware/Setup: Evaluated on 64 GPUs across multiple model configurations (e.g., DeepSeek-R1-Distill-Qwen-7B, Qwen3-30B-A3B, Qwen3-32B). Compared against the SOTA veRL and veRL+EPLB systems using PPO with datasets like DAPO-Math-17k and CodeForces.
• Performance: Achieves up to 1.45× speedup over state-of-the-art (SOTA) RL post-training systems (veRL) and up to 1.22× over systems enhanced with a step-level load balancer (veRL+EPLB).
• Granularity: Demonstrates that micro-step-level load balancing is necessary and superior to traditional step-level approaches in RL post-training.

Introduction and Theoretical Foundation

• Background: Focuses on the challenge of MoE load imbalance, which is exacerbated in RL post-training due to dynamic and imbalanced data (rollout batches) versus the more static data in pre-training.
• Key Observation: In RL post-training, routing decisions for a micro-batch are determined during the rollout stage and can be logged. This exposes precise future routing information, making load balancing a problem of planning with foreseeable routing, rather than prediction.
• Complexity: The load balancing problem at the micro-step level is NP-hard to solve optimally.

Micro-step-level Load Balancing

• Foreseeability: Routing decisions are exposed before the recompute and policy-update stages, enabling proactive planning. The target is to plan expert placement and token assignment for each micro-step, minimizing the max load across GPUs.
• Balancing Objective: The optimization minimizes the maximum rank load, including computation (from token assignments) and communication time (from inter-machine traffic), rather than merely balancing tokens per expert.

Hierarchical Planner

• Given an NP-hard problem, ForeMoE decomposes it into three tractable stages:
  1. Base Placement: Generates a stable expert placement per GPU using the aggregate token demand over the entire step.
  2. Intra-step Adjustment (Micro-step Level): For each micro-step, plans expert relocation and replication to adjust the base placement, mitigating transient load spikes without causing excessive transfer overhead.
  3. Token Assignment: Solves a linear program (LP) for optimal token routing to minimize the maximum rank load and communication cost.

Transfer Engine

• ForeMoE supports two complementary expert transfer paths:
  • CPU-assisted transfer: Moves experts via CPU memory, useful when GPU-GPU bandwidth is underutilized.
  • GPU-direct transfer: Transfers experts directly between GPUs, exploiting high-speed inter-GPU interconnects (e.g., NVLink) for lower latency.
• The system overlaps communication with computation by pipelining per-layer expert transfers during both the forward and backward passes. It uses a four-stage pipeline for CPU transfers (copy to CPU, H2D, D2H, copy from CPU) and a three-stage for GPU-direct transfers.

Theoretical Foundations

The load balancing problem is formulated as an NP-hard optimization. The paper defines the decision problem with a cost function C(π, α) = T_compute + T_comm and proves its NP-hardness theorem, establishing the theoretical need for a hierarchical planner.

The complexity of the problem motivates the hierarchical decomposition (base placement, micro-step adjustments, token assignment).

Conclusion

The paper presents ForeMoE, a micro-step-level load-balancing system for MoE during RL post-training. By exploiting the foreseeability of routing decisions, ForeMoE implements a two-level hierarchical planner to jointly optimize expert placement and token assignment, and a transfer engine with complementary CPU-assisted and GPU-direct paths. Evaluated on up to 64 GPUs, ForeMoE achieves up to 1.45× end-to-end speedup over state-of-the-art RL post-training systems.

Key Results

• Ablation Study: Confirms that all components—load balancing, expert transfer, and the hierarchical planner—contribute to the overall performance.
• Comparison with EPLB: ForeMoE significantly outperforms the step-level baseline (EPLB), highlighting the benefits of micro-step-level planning in dynamic RL environments.
• Ablation breakdown: Expert replication + Token assignment are more impactful in recompute; Expert relocation is key in policy update due to correctness constraints from micro-batching.
• Scalability: The system incurs low planning overhead (≈0.02% of step time) and supports large-scale models (up to 256 experts).
• Future Work: Extending ForeMoE to multi-node scenarios and integrating with sequence-level scheduling for further gains.

Theoretical Foundation

The proposed architecture is based on the observation that in RL post-training, unlike pre-training, routing decisions are predictable.

• Formalized Load Balancing Problem: Formulated as assigning experts to GPUs and tokens to experts to minimize makespan. This is NP-hard.
• Decomposition: The NP-hard problem is decomposed into tractable sub-problems: determining the expert-to-GPU placement, and optimizing the token-to-expert assignment for a given placement.

Key Insights

1. Foreseeability of Routing: RL algorithms like PPO and DAPO record routing decisions during rollout. The recompute stage can use this exact routing data, while the policy update stage can use the expected values from replay buffers. This makes the routing matrix fully known before execution.
2. Step-level Stability with Micro-step-level Variability: Long-term aggregate expert loads are stable, but individual micro-batches (the granularity of forward/backward passes) can be highly skewed. This insight pushes load balancing down to the micro-step level.

---

Introduction

• The Problem: Large Language Models (LLMs) are increasingly post-trained with Reinforcement Learning (RL) to improve reasoning and align with human preferences. This process is very computationally expensive, involving multiple training iterations of rollouts, reward scoring, and policy updates.
• The Bottleneck: MoE (Mixture-of-Experts) models are the dominant architecture for LLMs, but they suffer from load imbalance that causes significant efficiency bottlenecks. In RL post-training, this issue is exacerbated by dynamic prompt distributions.
• Current State-of-the-Art: Existing schedulers and memory optimizations for RL post-training are orthogonal to expert load balancing. Current load balancers (e.g., EPLB) operate at the pre-training step granularity and fail to leverage the foreseeability of routing decisions in RL.
• Key Opportunity: In RL post-training, routing decisions for a micro-batch are made during the rollout stage. In the subsequent recompute and policy-update stages, these decisions are known in advance (fixed), enabling a foreseeable routing paradigm.

Architecture and Key Components

• Micro-step-level Load Balancing: Leverages the foreseeability of routing in RL post-training to perform load balancing at the micro-step level, adapting to dynamic load variations within a training step.
• Hierarchical (Two-Level) Planner: Decomposes the NP-hard load-balancing problem (joint expert placement and token assignment) into:
  1. A base placement solver that runs at the step level.
  2. A token assignment planner that operates at the micro-step level.
• Dual-path Transfer Engine: Uses complementary hardware paths for expert transfer:
  • CPU-assisted transfer: For moving experts via the CPU (e.g., during the recompute stage) to overlap data movement with computation.
  • GPU-direct transfer: For maintaining low-latency expert swapping in subsequent stages (e.g., policy update), leveraging high-bandwidth GPU interconnects.

Conclusion

The paper addresses the critical challenge of load imbalance in MoE models during RL post-training. It identifies and characterizes the unique load characteristics where step-level expert loads are stable but skewed, while micro-step-level loads are highly variableqb. ForeMoE leverages the foreseeability of routing decisions in RL post-training to perform fine-grained, micro-step-level load balancing, decomposing the NP-hard problem into tractable sub-components via a hierarchical planner Matrix. It further designs a transfer engine that concurrently uses CPU-assisted and GPU-direct paths to handle expert transfers efficiently, hiding transfer overheads. The system achieves up to 1.45× end-to-end speedup over state-of-the-art RL post-training systems, demonstrating the significant benefits of exploiting routing foreseeability in this setting.