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

> ReLibra exploits routing replay in MoE RL training to decompose load balancing by timescale, achieving up to 1.58× throughput gains over Megatron-LM with negligible overhead.

- **Source:** [arXiv](https://arxiv.org/abs/2605.08639)
- **Published:** 2026-08-22
- **Permalink:** https://picx.dev/p/1TL5DI
- **Whiteboard:** https://picx.dev/p/1TL5DI/image

## Summary

## 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

| Metric | Improvement |
|--------|-------------|
| vs. Megatron-LM | 1.21–1.58× |
| vs. PopFetcher | 1.17–1.27× |
| vs. EPLB+ (oracle loads) | 1.06–1.21× |
| vs. LPLB+ (oracle loads) | 1.06–1.21× |
| vs. Balanced (ideal) | 90%–94% |

### Ablation Study

| Idx | Method | Reasoning Throughput | Mixed Throughput |
|-----|--------|---------------------|------------------|
| 1 | Megatron-LM | 1.0 | 1.0 |
| 2 | + inter-batch reordering | 1.12 (+12%) | 1.15 (+15%) |
| 3 | + intra-batch replication | 1.46 (+34%) | 1.39 (+24%) |

### 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

| Model | Replica Buffer Memory | Model Parameters | Ratio |
|-------|----------------------|------------------|-------|
| Qwen3-30B-A3B | 18 MiB | 6.26 GiB | 0.28% |
| GLM4.5-106B-A12B | 22 MiB | 16.64 GiB | 0.13% |
| Qwen3-235B-A22B | 72 MiB | 4.35 GiB | 1.62% |

- 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.

---

_Markdown view of https://picx.dev/p/1TL5DI, served by PicX — AI-generated visual whiteboard summaries of research papers._