# UltraEP: Unleash MoE Training and Inference on Rack-Scale Nodes with Near-Optimal Load Balancing

> UltraEP achieves 94.6% of ideal throughput in MoE training by using exact-load, real-time expert replication and rerouting, outperforming prediction-based balancers by 1.42x.

- **Source:** [arXiv](https://arxiv.org/abs/2606.04101)
- **Published:** 2026-08-18
- **Permalink:** https://picx.dev/p/2SbyCt
- **Whiteboard:** https://picx.dev/p/2SbyCt/image

## Summary

# UltraEP: Exact-Load Real-Time Load Balancing for MoE on Rack-Scale Nodes

## Summary (Overview)

- **UltraEP** is the first exact-load, real-time load balancer for large-scale expert parallelism (EP) in MoE training and serving prefill on rack-scale nodes (RSNs), rebalancing every microbatch and layer on the critical path.
- Achieves **94.6% of force-balanced ideal throughput in training** and **93.9% in serving prefill** on average across models from 106B to 671B parameters on up to 256 GPUs.
- Introduces a **quota-driven planner** that jointly optimizes expert replication and token rerouting using exact post-gating load, with a threshold-based binary search feasibility oracle.
- Implements **RSN-native communication** via persistent tile streaming and chunk-streaming relay trees for hot-expert fan-out, achieving 3.1×–5.5× speedup over mainstream communication backends.
- Reduces final inter-rank imbalance from 1.30–4.01 to **1.01–1.04**, with stable production performance exceeding 92% of ideal throughput while preserving training convergence.

## Introduction and Theoretical Foundation

### Motivation

Large-scale expert parallelism (EP) is essential for training and serving frontier MoE models but amplifies expert load imbalance. The paper identifies three key challenges:

1. **Compute stragglers**: Tokens distribute unevenly across experts due to routing dynamics
2. **Token all-to-all bottlenecks**: Communication hotspots on overloaded devices
3. **Activation-memory spikes**: Memory pressure on overloaded devices

### Limitations of Existing Approaches

Existing balancers (e.g., EPLB) rely on **historical load prediction** and periodic rebalancing, which becomes unreliable with non-stationary load patterns. The paper demonstrates that expert popularity shifts sharply across microbatches, layers, and data domains, making stale predictions produce inaccurate placements.

### Rack-Scale Node (RSN) Opportunity

The emergence of RSNs fundamentally changes the design space by extending scale-up connectivity across 64+ GPUs in a full rack, making hot-path balancing physically viable. Key properties:
- Scale-up bandwidth: hundreds of GB/s per GPU (8–10× scale-out)
- Load/store-style memory semantics
- EP groups can be contained within one RSN

## Methodology

### System Architecture

UltraEP implements a **replication-only** balancing strategy with several key design decisions:

**1. Expert Layout and Memory Management:**
- Logical experts (model-defined) map to physical experts (replicas on ranks)
- Fixed layout: one main instance + zero or more redundant replicas per logical expert
- Cross-layer buffer reuse: redundant slots share weight/gradient buffers across layers, reducing per-slot memory from 3.3GB to 36MB for weights

**2. Quota-Driven Planning** (Algorithm 1):

The planner solves a joint replication-reroute problem:

$$T_{\text{solve\_rep}}^{fwd} + \max(T_{\text{reroute}}^{fwd}, T_{w\_distr}^{fwd}) + T_{\text{tok\_a2a}}^{fwd} + T_{\text{moe}}^{fwd}$$

Key components:
- **Threshold formulation**: Binary search for the smallest load threshold $\tau$ such that all ranks can be brought below it via replication
- **Feasibility oracle**: Greedy transfer of excess load from overloaded ranks to admissible ranks with slack
- **Quota construction**: Each replica carries a minimum useful quota $u_{\min} = 1024$ tokens

**3. Reroute with Locality:**

$$q_{t,e,t} \leftarrow \min(\lambda_{t,e}, u_{e,t}) \quad \text{(local quota consumption first)}$$

### RSN-Native Communication

**Persistent Tile Streaming:**
- Expert weights/gradients divided into fixed-size tiles
- Persistent kernel with double-buffered shared memory
- Overlap-aware footprint: configurable SM residency for forward vs. backward paths

**Chunk-Streaming Relay for Hotspot Fan-Out:**
- Two-stage relay for experts with replica count > 4
- Relay frontier at approximately $\sqrt{|\mathcal{H}(e)| - 1}$
- Chunk-level pipelining without global barriers
- Load-aware relay scheduling to balance outgoing traffic

### GPU-Native Solving

- Fully on-device implementation (no CPU synchronization)
- Warp-level parallelism with shared memory staging
- Multiple threshold probes evaluated across warps

## Empirical Validation / Results

### Setup

| Model | Experts (Top-k) | Parallelism (Train \| Serve) | $N_{\text{slot}}$ |
|-------|-----------------|------------------------------|--------------------|
| GLM4.5-106B-A12B | 128 (8) | EP64-DP2 \| - | 2 |
| Qwen3-235B-A22B | 128 (8) | EP64-DP4 \| EP64 | 2 |
| GLM4.7-358B-A32B | 160 (8) | - \| EP40 | 4 |
| DeepSeek-V3-671B-A37B | 256 (8) | EP64-PP4 \| - | 2 |

### Key Results

**Training Throughput Improvement:**
- **1.42×** over Megatron-LM (average)
- **1.56×** over SGLang for serving prefill
- 93.9–94.6% of force-balanced ideal throughput

**Balancing Quality:**
- Post-balancing imbalance: **1.01–1.04** (vs. 1.30–4.01 without)
- 57.9% fewer redundant slots than EPLB+
- 3.9% traffic reduction via locality

**Communication Speedup:**
- 3.1×–5.5× over torch.distributed and DeepEP
- Relay provides additional 1.3×–1.8× for large fan-out

### Comparison with Baselines

| Metric | EPLB+ | UltraEP |
|--------|-------|---------|
| Result Imbalance | 1.19 | **1.03** |
| Solving Time (ms) | 0.153 | **0.111** |
| Total Replicas | 107 | **45** |
| Max Fan-out | 8.5 | **6.8** |
| In-flight Token Ratio | 99.9% | **96.0%** |

## Theoretical and Practical Implications

### System Design Insights

1. **Exact load beats prediction**: Reacting to realized load with efficient planning (0.111ms solve time) outperforms history-based prediction that can be arbitrarily wrong under non-stationary loads.

2. **Joint planning is crucial**: Directly optimizing post-reroute load (quota) is superior to decoupled placement-then-reroute approaches, using 57.9% fewer resources.

3. **RSN-native communication is essential**: Purpose-built communication kernels for dynamic expert traffic achieve near-constant latency (0.28ms) even with increasing fan-out, while generic backends scale linearly.

### Production Impact

- **Memory efficiency**: Reduces MoE activation peak memory by up to 11× in serving, lowering OOM risk
- **Convergence preservation**: Loss curve follows expected pretraining trajectory (replication-only, no reordering)
- **Scalability**: Stable performance >92% of ideal throughput in production across multiple racks

## Conclusion

UltraEP demonstrates that **exact-load, real-time expert balancing is both feasible and near-optimal** on RSNs. The key insight is that with sufficient communication bandwidth (RSN scale-up), the balancing problem shifts from "how to predict and hide" to "how to solve and communicate fast enough."

### Future Directions

1. **Reinforcement learning pipelines**: Extending the same abstraction to RL training that alternates training and inference
2. **Decode-phase optimization**: Currently focused on prefill; decode presents different challenges due to memory-bound characteristics
3. **Cross-RSN balancing**: Handling models that span multiple RSNs where bandwidth is more constrained
4. **Adaptive slot allocation**: Dynamic $N_{\text{slot}}$ adjustment based on observed load patterns

The work establishes that with appropriate system co-design, the theoretical benefits of load balancing can be realized in practice, achieving near-force-balanced performance in production MoE deployments.

---

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