Summary (Overview)

• Core finding: Scaling expert parallelism (EP) does not reduce routing imbalance in Mixture-of-Experts (MoE) models—the straggler is intrinsic to the routing decision, not the expert-to-rank mapping.
• Benchmark invalidity: Mock-token benchmarks overestimate routing imbalance by up to 2.35×, and this error is a constant level offset that persists across a 32× batch-size sweep, not a trend.
• Architecture taxonomy: Five architectures separate into two bands—data-resilient (MHA, Mamba-2) where real text drives routing near-uniform, and persistently concentrated (MLA, GDN) where Gini stays high regardless of input; GQA is intermediate.
• Methodology: A 5×6 factorial study (5 architectures × 6 data conditions) on H100 clusters with matched-window EP scans and controlled batch-size sweeps, using released checkpoints with as-deployed routing configurations.
• Key implication: The architecture class, not EP degree or mock-data profile, is the correct workload input for AlltoAll-aware interconnect and dispatch design.

Introduction and Theoretical Foundation

The paper addresses a critical bottleneck in frontier MoE training: roughly half of forward-pass time is spent in two AlltoAll calls per MoE block. Four families of system-level responses have emerged to address this:

1. Predictive placement exploiting routing locality across iterations
2. Adaptive expert relayout responding to measured load
3. Hierarchical collectives aggregating AlltoAll traffic in stages
4. EP-aware topology redesigning networks around EP traffic

These approaches share two untested assumptions:

• Assumption A (correctability): Routing imbalance is a workload property a system layer can reduce
• Assumption B (benchmark validity): Mock/random token benchmarks faithfully represent production routing

The paper's core theoretical foundation rests on the AlltoAll bottleneck mechanics. With expert parallelism across $P$ ranks, token dispatch creates a send counts matrix $S \in \mathbb{N}^{P \times P}$ where $S_{ij}$ is the number of tokens rank $i$ sends to rank $j$. The collective completes only when every rank has received all incoming tokens, so completion time tracks the maximum column sum, not the mean.

The paper introduces three scalar metrics to characterize imbalance:

Per-expert max/mean ratio: $\max_{\ell} e_{\ell} / \bar{e}$ — the primary statistic for EP scaling questions, capturing routing decisions in a form EP cannot influence.

Gini coefficient of per-rank receive counts, with sorted $c_{(1)} \leq \cdots \leq c_{(P)}$:

$$G = \frac{\sum_{j=1}^{P} (2j - P - 1) c_{(j)}}{P \sum_{j} c_{j}} \tag{1}$$

ranging from 0 (perfect balance) to 1 (one rank absorbs all).

Fitted Dirichlet concentration α: treats per-rank receive proportions as a draw from a symmetric Dirichlet, measuring how tightly the distribution clusters around uniform. α → ∞ means concentrated at uniform; α → 0 means concentrated at a single-rank vertex.

Methodology

Hardware setup: All measurements on the Helma supercomputer at NHR@FAU: four H100 GPUs per node (94 GB HBM2e each), NVLink intra-node, NDR200 InfiniBand inter-node. Software: PyTorch with NVIDIA NeMo 26.02 and Megatron-Core 0.16.1, sequence length 4096, bf16 precision.

Five architectures tested (Table I):

Six data conditions forming a gradient isolating independent properties:

1. Mock – uniform-random token IDs (de facto benchmark default)
2. Shuffled – real wikitext tokens, globally permuted (preserves frequency, destroys sequence)
3. Remapped – real wikitext sequences with permuted vocabulary IDs (preserves sequence, destroys embeddings)
4. Romansh – real text in Rumantsch Grischun (unfamiliar language)
5. Opus – post-cutoff English reasoning traces (familiar language, unseen content)
6. Wikitext – English Wikipedia (familiar language, likely seen content)

Experimental design:

• Factorial 5×6 block: EP=16, GBS=16, 5 warmup + 10 measurement iterations
• EP scan: GBS=32, 50 warmup + 200 measurement, EP ∈ {4, 8, 16, 32} for MLA/MHA; EP ∈ {8, 16, 32} for others (memory constraints)
• Batch-size sweep: GBS ∈ {64, 256, 1024, 2048} under mock and wikitext

Instrumentation: Logs per-dispatch input_splits from MoEAlltoAllTokenDispatcher on every EP rank at every MoE layer, recovering the full S matrix. Timing via CUDA events with system P99 as max over ranks.

Empirical Validation / Results

Q1: EP Scaling Does Not Reduce the Straggler

The EP scan shows five flat lines (≤5% variation in max/mean per architecture):

• MLA: 5.0% variation across EP={4, 8, 16, 32}
• MHA: 2.8%
• GQA: 4.4% across EP={8, 16, 32}
• Mamba-2: 4.4%
• GDN: 0.3%

The flatness holds despite the EP axis spanning an ~18× per-GPU bandwidth transition (NVLink-only at EP=4 to IB-dominated at EP=32). Absolute dispatch P99 does scale (~15× increase), but the skew multiplier is EP-invariant. The interpretation is mechanical: EP scaling changes which rank holds which expert, not which experts the model sends tokens to.

Metric validation: across the 30 factorial cells at EP=16, per-cell Gini predicts dispatch P99 with $r \geq 0.76$ within each architecture ($r \geq 0.99$ for P50).

Q2: Mock Data Overestimates Imbalance at Every Batch Size

The mock-to-wikitext Gini ratio ranges from 1.02× (GDN) to 2.35× (Mamba-2). Key controls:

• Shuffled condition routes as badly as or worse than mock for every architecture (Mamba-2: Gini 0.491, 39% worse than mock)
• Remapped condition degrades routing to mock-tier, showing the router needs correct learned embeddings
• Zipf-skewed synthetic tokens route worse than uniform mock for every architecture (overestimate reaches 3.9× against wikitext)

Batch-size invariance: Per-rank Gini varies by less than 1% (MLA) and ~3% (MHA) under mock across a 32× GBS sweep, with seed-to-seed spread (std ≤ 0.004) of the same order as total sweep variation. A single-sample measurement can fabricate apparent trends (+17% to +40%) that hundreds of iterations refute.

Q3: Two Emergent Architecture Classes

Data-resilient class (MHA, Mamba-2): Lowest absolute Gini under wikitext (0.105 and 0.150), largest gain from real data (2.24× and 2.35× improvement from mock).

Persistently concentrated class (MLA, GDN): Gini stays above 0.24 under English real text, above 0.21 everywhere; smallest improvement (1.56× and 1.02×).

Intermediate (GQA): Wikitext Gini (0.240) in MLA's band, but mock-to-wikitext improvement (1.43×) closer to MHA's.

Per-layer analysis reveals distinct depth profiles:

• MHA: shallow U-shape (layer-0 Gini 0.17, dipping to 0.12, rising to 0.30)
• Mamba-2: single localized spike of 0.63 at depth fraction 0.45 (preceding an attention layer), otherwise ≤0.40
• MLA: flat-high, fluctuating 0.35–0.50
• GQA: rises sharply to ~0.59 by depth 0.20, stays 0.55–0.67
• GDN: oscillates 0.40–0.63 with no clear direction

Mechanism findings:

• Router swap experiment shows routing Gini follows the model body, not the router (MHA hidden states route at Gini ~0.27 regardless of router; MLA hidden states at ~0.34 either way)
• Random-init experiments show the two-class taxonomy emerges from pretraining, not the forward pass
• Mamba-2 shows a distinctive inversion: worst under mock (plateau 0.15–0.30, no descent) but best under wikitext (descends to 0.047), consistent with an SSM amplifying sequential statistics

Theoretical and Practical Implications

For the four system families:

• Predictive placement: Temporal stability analysis finds $r > 0.7$ between consecutive steps for MHA, GQA, MLA, Mamba-2 (GQA reaching $r = 0.93$, lag-20 at 0.91), but GDN collapses to $r \approx 0$ within five steps
• Adaptive relayout: EP-invariance implies return-on-cost should be evaluated against an absolute target, with the persistent class as natural beneficiary
• Hierarchical AlltoAll: Persistent class produces substantially more on-node concentration than resilient class
• EP-aware topology: Adds bandwidth helping all architectures uniformly, but workload models need calibration on real data

For practitioners:

1. Report both mock and real-data conditions in benchmarks
2. Treat the three regimes separately:
   • Data-resilient pair: uniform-load models are appropriate
   • Persistent pair: expect Gini > 0.24 under English real text at every batch size and EP; locality-aware dispatch has more to gain
   • GQA: size for inter-node concentration, prime candidate for predictive placement

Key limitations:

• Scale: factorial uses 16 GPUs, EP scan reaches 32 GPUs; production uses 256–4096 GPUs
• All models use auxiliary-loss-based balancing; aux-loss-free routing not tested
• Released checkpoints profiled over short windows; continued-training adaptation changes the gap

Conclusion

Routing imbalance in MoE expert parallelism is not a scaling problem but a model-and-data property. The synthetic-token benchmarks used to design MoE-aware interconnect proposals overestimate imbalance by a constant factor that persists across batch sizes, and skewing synthetic distributions toward realism (Zipf) widens rather than closes the gap. The five architectures separate into data-resilient (MHA, Mamba-2), persistently concentrated (MLA, GDN), and intermediate (GQA) classes with qualitatively different load profiles for interconnect and dispatch design.

Future work includes: production-scale expert parallelism (EP=64–320), routing dynamics under continued training, trace-driven network simulation replaying measured send-counts matrices, and A/B testing of locality-aware dispatch for the persistent class.