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_{\mathrm{solve\_rep}}^{fwd} + \max(T_{\text{reroute}}^{fwd}, T_{w\_distr}^{fwd}) + T_{\mathrm{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

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

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.