Summary (Overview)

• ReLibra is an MoE (Mixture-of-Experts) 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.
• The key insight is to place two load-balancing mechanisms at different timescales: inter-batch expert reordering for coarse-grained cross-node balancing, and intra-batch expert replication for fine-grained local adaptation, matching their communication patterns to hierarchical network bandwidths.
• ReLibra decomposes the NP-hard load-balancing problem into two tractable sub-problems, solved via a swap-based simulated annealing algorithm (inter-batch) and an incremental greedy heuristic with MILP formulation (intra-batch).
• A layer-shared replica buffer stores replicas for only one layer at a time, reducing memory overhead by 1–2 orders of magnitude, with synchronization overlapped with attention computation.
• Experiments show ReLibra improves training throughput by up to 1.6× over Megatron-LM and up to 1.2× over EPLB+ (even with oracle loads), achieving 90%–94% of an idealized balanced baseline.

Introduction and Theoretical Foundation

Background: MoE Training and RL Workflow

MoE (Mixture-of-Experts) architectures scale LLMs by sparsely routing tokens to a small set of experts, increasing model capacity without proportionally increasing computation. However, MoE training suffers from load imbalance—dynamic routing can overload GPUs hosting "hot" experts under expert parallelism (EP).

The paper identifies a critical distinction between pretraining and RL post-training:

• MoE pretraining exhibits relatively stable imbalance patterns due to web-scale diverse corpora and algorithmic load-balancing mechanisms (auxiliary losses, routing bias).
• MoE RL training shows fluctuating imbalance patterns: hot experts shift substantially from one micro-batch to the next because (1) RL targets narrower domains (math, coding, instruction following), and (2) load-balancing mechanisms are removed during RL training.

The Opportunity: Routing Replay

RL consists of two stages:

1. Rollout stage: generates responses via auto-regressive decoding
2. Training stage: updates the model on generated samples

Since both stages process the same token sequence with the same MoE parameters, they produce the same routing scores and top-k expert selections. Thus, token-to-expert routing decisions are known before training starts—this is "routing replay."

The Challenge: Practical Load-Balancing Planning

Even with exact routing known, finding optimal load-balancing plans is computationally intractable:

• Without replication and memory constraints, the problem reduces to makespan minimization on identical machines ($P||C_{\text{max}}$)—NP-hard
• With replication and capacity constraints, it contains replica placement decisions similar to the facility location problem

Key Insight: Timescale Decomposition

Under cross-node EP:

• Expert reordering is a global operation requiring coordinated placement changes across the EP group—poorly suited for micro-batch-level fluctuations due to expensive cross-node coordination
• Expert replication can be restricted within a node using high-bandwidth intra-node links—well-suited for micro-batch timescale

Therefore: expert reordering at inter-batch timescale, expert replication at intra-batch timescale.

Methodology

System Architecture

ReLibra has two planes:

Control Plane: Maintains global view, receives routing information from rollout, solves load-balancing problems (overlapped with post-rollout processing).

Data Plane:

• Routing Capturer: Records routed experts for each token during rollout generation
• Load-Balancing Executor: Executes reordering and replication plans, manages the layer-shared replica buffer

Inter-Batch Expert Reordering

Problem Formulation: For each MoE layer, ReLibra jointly optimizes computation and communication loads. The total load of GPU $g$ is:

$$L_g^{\text{comp}} = \sum_{e \in E_g} l_e = \sum_{e \in E_g} \sum_{j \in \mathcal{G}} x_{j,e} \tag{1}$$

where $x_{j,e}$ is the number of tokens on GPU $j$ routed to expert $e$.

Communication load accounts for the rail-optimized topology with three patterns: intra-node NVLink, inter-node same-rail RDMA, and inter-node cross-rail (NVLink relay + RDMA). The MoE execution time is:

$$T_{\text{MoE}} = \max_{g \in \mathcal{G}} T_g^{\text{comp}} + \max_{g \in \mathcal{G}} T_g^{\text{comm}} \tag{6}$$

Algorithm: A swap-based simulated annealing algorithm:

• Initial state from greedy LPT (Longest Processing Time) heuristic
• Repeatedly swaps randomly chosen experts with incremental state updates ($O(|\mathcal{G}|)$ complexity per swap)
• Objective smoothing via log-sum-exp (LSE) surrogate to escape plateaus: $\text{LSE}_\beta(z_1,\dots,z_n) = \frac{1}{\beta}\log\sum_{k=1}^{n} e^{\beta z_k}$
• Multiple parallel annealing runs with different seeds
• Second round of optimization for data locality (sample-to-GPU assignment)

Intra-Batch Expert Replication

Layer-Shared Replica Buffer: Each GPU allocates only a small number of replica slots, stores replicas for one layer at a time, and reuses the buffer across all layers. This reduces memory overhead from $Lr$ replica slots (naive) to just $r$ slots.

MILP Formulation: For each expert $e$ in a micro-batch, the system jointly optimizes replica placement ($z_{e,g}$ binary variables) and token splitting ($y_{j,e,g}$ continuous variables):

$$\min_{y_{j,e,g} \in [0,1], z_{e,g} \in \{0,1\}} T_{\text{MoE}} = T^{\text{comp}}(y,z) + T^{\text{comm}}(y,z) \tag{7}$$

Subject to constraints:

• $\sum_{g \in \mathcal{G}(e)} y_{j,e,g} = 1$ (token fractions sum to 1)
• $y_{j,e,g} \leq z_{e,g}$ (coupling token splitting with placement)
• $z_{e,g} = 1$ for home GPUs
• $\sum_{e \in E \setminus E_g: g \in \mathcal{G}(e)} z_{e,g} \leq r$ (per-GPU replica buffer limit)

Incremental Greedy Heuristic: Avoids direct MILP solving by iteratively:

1. Identifying the bottleneck GPU (highest load)
2. Selecting the bottleneck expert on that GPU
3. Adding a replica to the least-loaded feasible GPU
4. Re-solving token splitting via linear programming

Empirical Validation / Results

Testbed and Models

• Hardware: 20-node cluster, 8 NVIDIA Hopper GPUs per node, rail-optimized InfiniBand (8×400 Gbps per node), 900 GB/s NVLink
• Models: Qwen3-30B-A3B, GLM4.5-106B-A12B, Qwen3-235B-A22B (all with 128 experts/layer, top-8 routing)
• Datasets: Reasoning (DAPO-Math-17k, GPQA), Instruction Following (IFBench), Coding (CodeForces), and Mixed

Overall Performance

Ablation Study

Case Study: Load Imbalance Reduction

• EP size study (Qwen3-30B-A3B): ReLibra maintains average skewness of 1.00–1.08 across EP sizes 8–64 (1.0 = perfectly balanced), while Megatron-LM degrades significantly with larger EP sizes
• Across datasets (Qwen3-235B-A22B): ReLibra achieves average skewness of 1.02–1.07

System Overhead

• Solving time fully hidden behind post-rollout processing
• Expert reordering overhead: only 1.4%–2.1% of training batch time
• Replica synchronization fully overlapped with attention computation

Theoretical and Practical Implications

Theoretical Contributions

1. Problem decomposition: ReLibra demonstrates that the intractable joint optimization of expert reordering and replication can be decomposed by timescale, matching communication patterns to network hierarchy—making the problem tractable while achieving near-optimal results.

2. Routing replay as a paradigm: The paper formalizes routing replay as a load-balancing opportunity unique to RL training, contrasting with the prediction-based approaches used in pretraining.

3. Holistic replication planning: The MILP formulation jointly optimizes replica placement and token splitting, showing that these coupled decisions outperform optimizing either in isolation.

Practical Implications

1. Industrial relevance: The paper aligns with recent industrial practice (DeepSeek-V3.2) that has adopted routing replay for training stability, making ReLibra's approach immediately applicable.

2. Memory efficiency: The layer-shared replica buffer design makes micro-batch-level replication practical for deep MoE models, reducing memory overhead by 1–2 orders of magnitude.

3. Orthogonality: ReLibra is orthogonal to existing all-to-all optimizations (DeepEP, hierarchical communication) and can be combined with them.

Conclusion

ReLibra addresses fluctuating load imbalance in MoE RL training by exploiting routing replay. Key takeaways:

1. Timescale decomposition is the core insight: expert reordering at inter-batch timescale for cross-node balancing, expert replication at intra-batch timescale for local adaptation—matching communication patterns to network bandwidth hierarchy.

2. Routing replay transforms load balancing from prediction-based to knowledge-based, enabling proactive fine-grained decisions at micro-batch granularity.

3. Efficient algorithms (swap-based simulated annealing, incremental greedy heuristic) make the NP-hard problem tractable, with solving time hidden behind post-rollout processing.

4. Lightweight system design (layer-shared replica buffer, overlapped synchronization) achieves load balancing with negligible overhead (0.13%–1.62% additional memory).

Future directions include extending to off-policy RL with asynchronous rollout-training workflows, where routing replay may not directly apply but recent industrial practice suggests it can be adopted for training stability.