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

• URL: http://arxiv.org/abs/2606.04101v2
• Date: Fri, 05 Jun 2026 06:53:34 GMT
• Title: UltraEP: Unleash MoE Training and Inference on Rack-Scale Nodes with Near-Optimal Load Balancing
• Abstract summary: We present UltraEP, the first exact-load, real-time balancer for large-EP MoE training and serving prefill on rack-scale nodes (RSNs)<n>Built upon the extended scale-up connectivity of RSNs, UltraEP rebalances every microbatch and layer on critical paths.<n>Averaged across MoE models from 106B to 671B parameters in training and prefill, UltraEP achieves 94.3% of the force-balanced ideal throughput.
• Score: 28.287651382171692
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Large-scale expert parallelism (EP) is becoming pivotal for training and serving frontier MoE models, but it also amplifies device-level expert load imbalance into compute stragglers, token all-to-all bottlenecks, and activation-memory spikes. Existing balancers redistribute experts periodically based on historical load, which becomes unreliable for production deployments with non-stationary load patterns. We present UltraEP, the first exact-load, real-time balancer for large-EP MoE training and serving prefill on rack-scale nodes (RSNs). Built upon the extended scale-up connectivity of RSNs, UltraEP rebalances every microbatch and layer on critical paths, which requires nontrivial co-design of plan solving and expert replication communication to minimize exposed overhead. To this end, UltraEP eagerly reacts to post-gating load with efficient quota-driven planning, and executes the resulting irregular expert-state transfers with RSN-native persistent tile streaming and relay-based fan-out mitigation. Averaged across MoE models from 106B to 671B parameters in training and prefill, UltraEP achieves 94.3% of the force-balanced ideal throughput, delivering 1.49$\times$ improvement over non-balancing, while reducing the final inter-rank imbalance from 1.30$-$4.01 to 1.01$-$1.04. Additionally, we validate UltraEP's scalability and robustness in production MoE training with 2560 GPUs.

Related papers

• $φ$-Balancing for Mixture-of-Experts Training [20.061328126956028]
  Mixture-of-Experts (MoE) models rely on balanced expert utilization to fully realize their scalability.<n>We propose $$-balancing, a principled framework that directly targets population-level expert balance by minimizing a strictly convex, symmetric, and differentiable potential of the expected routing distribution.
  arXiv  Detail & Related papers  (2026-05-14T20:39:28Z)
• TurboGR: An Accelerated Training System for Large-Scale Generative Recommendation [9.645364292862624]
  Generative recommendation (GR) has emerged as a promising paradigm that replaces fragmented, scenario-specific architectures with unified Transformer-based models.<n> deploying GR at scale on Ascend NPUs faces fundamental system-level challenges.<n>We present model, an Ascend-affinity training system for generative recommendation.
  arXiv  Detail & Related papers  (2026-05-13T12:26:29Z)
• FEMBA on the Edge: Physiologically-Aware Pre-Training, Quantization, and Deployment of a Bidirectional Mamba EEG Foundation Model on an Ultra-low Power Microcontroller [13.961569964076451]
  We present FEMBA, a bidirectional Mamba architecture pre-trained on over 21,000 hours of EEG.<n>We introduce a novel Physiologically-Aware pre-training objective, consisting of a reconstruction with low-pass filtering.<n>We employ Quantization-Aware Training (QAT) to compress the model to 2-bit weights.
  arXiv  Detail & Related papers  (2026-03-18T13:23:13Z)
• A Replicate-and-Quantize Strategy for Plug-and-Play Load Balancing of Sparse Mixture-of-Experts LLMs [64.8510381475827]
  Sparse Mixture-of-Experts (SMoE) architectures are increasingly used to scale large language models efficiently.<n>SMoE models often suffer from severe load imbalance across experts, where a small subset of experts receives most tokens while others are underutilized.<n>We present a systematic analysis of expert routing during inference and identify three findings: (i) load imbalance persists and worsens with larger batch sizes, (ii) selection frequency does not reliably reflect expert importance, and (iii) overall expert workload and importance can be estimated using a small calibration set.
  arXiv  Detail & Related papers  (2026-02-23T15:11:16Z)
• ECO: Quantized Training without Full-Precision Master Weights [58.97082407934466]
  Error-Compensating (ECO) eliminates master weights by applying updates directly to quantized parameters.<n>We show that ECO converges to a constant-radius neighborhood of the optimum, while naive master-weight removal can incur an error that is inversely proportional to the learning rate.
  arXiv  Detail & Related papers  (2026-01-29T18:35:01Z)
• Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts [74.40169987564724]
  Expert parallelism (EP) is designed to scale MoE models by distributing experts across multiple devices.<n>Under extreme imbalance, EP can funnel a disproportionate number of tokens to a small number of experts, leading to compute- and memory-bound failures.<n>We propose Least-Loaded Expert Parallelism (LLEP), a novel EP algorithm that dynamically reroutes excess tokens and associated expert parameters from overloaded devices to underutilized ones.
  arXiv  Detail & Related papers  (2026-01-23T18:19:15Z)
• Controlled LLM Training on Spectral Sphere [76.60985966206746]
  We introduce the textbfSpectral Sphere algorithm (SSO), which enforces strict module-wise spectral constraints on both weights and their updates.<n>We observe significant practical stability benefits, including improved MoE router load balancing, suppressed outliers, and strictly bounded activations.
  arXiv  Detail & Related papers  (2026-01-13T09:59:47Z)
• MC#: Mixture Compressor for Mixture-of-Experts Large Models [86.64315380917827]
  Mixture-of-Experts (MoE) effectively scales large language models (LLMs) and vision-language models (VLMs) by increasing capacity through sparse activation.<n>We propose MC# (Mixture-Compressor-sharp), a framework that combines static quantization and dynamic expert pruning.
  arXiv  Detail & Related papers  (2025-10-13T03:12:46Z)
• PT$^2$-LLM: Post-Training Ternarization for Large Language Models [52.4629647715623]
  Large Language Models (LLMs) have shown impressive capabilities across diverse tasks, but their large memory and compute demands hinder deployment.<n>We propose PT$2$-LLM, a post-training ternarization framework tailored for LLMs.<n>At its core is an Asymmetric Ternary Quantizer equipped with a two-stage refinement pipeline.
  arXiv  Detail & Related papers  (2025-09-27T03:01:48Z)
• Latent Prototype Routing: Achieving Near-Perfect Load Balancing in Mixture-of-Experts [0.0]
  Latent Prototype Routing (LPR) is a novel routing framework that promotes balanced expert utilization without compromising downstream performance.<n>LPR reduces the Gini coefficient of expert load from 0.70 to 0.035 on average, improves the min-max expert load ratio from 1e-6 to 0.70, achieving near-perfect load balancing.
  arXiv  Detail & Related papers  (2025-06-26T14:41:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.