MoE Pretraining: Parallelism Layout Changes the Model, Not Quality

Summary (Overview)

  • This paper investigates whether the parallelism layout (data parallel / expert parallel configuration) changes the trained MoE model, even when seed, data, and hyperparameters are held fixed.
  • Across 80 training runs, the author finds that layout does not change model quality (final validation loss stays within seed-to-seed variation), but it does change the model itself — predictions differ about 80% as much as a new random seed would make them, with over 90% of routing decisions differing.
  • The root cause is identified as near-ties in the router: routing decisions that flip between runs are almost exactly those where the chosen expert and the best unchosen expert had nearly equal scores (AUC 0.97–1.00 in early training).
  • Fine-grained experts make routing less reproducible: going from 8 to 64 experts raises near-tie decisions from 2.7% to 23.3% and routing differences from 68% to 100%, though output drift stays constant.
  • Timing of layout switches matters: switching at 84% of pretraining barely changes the final model, while switching at 10% (just after warm-up) changes it substantially.

Introduction and Theoretical Foundation

Background

Mixture-of-Experts (MoE) models replace one large feed-forward block with many smaller "experts" and a router that sends each token to a few of them (Shazeer et al., 2017; Lepikhin et al., 2021; Fedus et al., 2022). Training requires many machines, with work split via:

  • Data parallelism (DP): each machine processes part of the batch
  • Expert parallelism (EP): each machine holds part of the experts, with token exchange between machines

The choice of DP and EP sizes is called the layout.

Key Insight

In practice, the layout is not fixed—jobs are preempted and resumed on different GPU counts, clusters resize, and ablations run on new setups. The usual assumption is that this is harmless, but:

  1. Numerical effects: Computers add numbers with finite precision, so (a+b)+c≠a+(b+c)(a + b) + c \neq a + (b + c), and changing layout changes the order of gradient and expert-output additions.
  2. Definitional effects: MoE rules (load-balancing statistics, expert capacity, loss-free balancing counts) are often computed per machine rather than over the whole batch, meaning a new layout changes what is computed, not just its rounding.

Research Questions

  1. Does the parallelism layout change the trained MoE model?
  2. If so, why?
  3. Does it get worse as experts become more fine-grained?

Measurement Framework

The author compares runs against three yardsticks: a new random seed, a new data order, and a 1-ULP (one unit in the last place) change to a single weight. The main difference measure is the KL divergence between predicted next-token distributions on 16,384 held-out tokens, divided by the KL between two seed-differing runs ("vs. own seed").


Methodology

Models and Data

ItemSetting
Base MoE6 layers, width 256, 4 heads, RoPE, RMSNorm, QK-Norm, no QKV bias, tied embeddings
Experts16 SwiGLU experts (width 256), top-4 routing, no shared expert; ~25M total / ~10M active params
Dense controlSame model, one SwiGLU block of width 1024 (same active size)
DataFineWeb-Edu; 16,384-token byte-level BPE; 20M training tokens (1,221 steps), global batch 32 × 512 tokens
OptimizerAdamW, peak LR 10−310^{-3}, 100 warm-up steps, weight decay 0.1, β=(0.9,0.95)\beta = (0.9, 0.95), clip 1.0
Precisionfp32 on CPU (Azure E16s_v3), fp16 on GPU (NVIDIA T4)

Layouts Tested

Nine DP/EP layouts were evaluated: (1,1), (2,1), (2,2), (4,1), (4,2), (4,4), (8,1), (8,2), (8,4), using real multi-process PyTorch with gloo communication.

Correctness Checks

Before training, the author verified that:

  • Every layout computes the same gradients as a single process (relative error < 10−410^{-4} in fp32)
  • Two runs of the same layout are bit-identical
  • Checkpoints move between layouts without change

Two Kinds of Layout Effects

  • Numerical effect: only the order of additions changes
  • Definitional effect: a per-machine rule (load balancing, capacity, loss-free bias) is switched on

Empirical Validation / Results

Finding 1: Same Quality, Different Model

ComparisonKL at step 8KL at endvs. own seedTop-1 agreeRouting differsLoss barrierLoss change
MoE: different seed1.0×10−11.0 \times 10^{-1}0.3501.000.664-3.538+0.023
MoE: different layout (8)4.1×10−94.1 \times 10^{-9}0.2800.800.68992%0.417+0.005
MoE: 1-ULP nudge (3)3.6×10−93.6 \times 10^{-9}0.2830.810.68593%0.432+0.016
MoE: thread count6.6×10−96.6 \times 10^{-9}0.2770.790.68992%-+0.006
Dense: different layout (2)00.0430.120.862---0.002
Dense: 1-ULP nudge (2)00.0480.130.859-0.000-0.002

Key observations:

  • The difference starts tiny (4×10−94 \times 10^{-9} at step 8) but grows to 0.80× the seed level by end of training
  • All eight layouts, 1-ULP nudges, and thread-count changes follow one curve from step 32 onward
  • Final validation losses (4.576–4.610) fall inside the range of new-seed runs (4.582–4.626)

Finding 2: Per-Machine Rules Start Large but End the Same

Per-machine rules create differences ~400× larger at step 8 (1.7×10−61.7 \times 10^{-6} vs. 4×10−94 \times 10^{-9}), but by step 64 they are indistinguishable from pure layout changes, ending at 0.282 vs. 0.280.

Finding 3: The Cause is Near-Ties in the Router

The author defines a routing decision's margin as the gap between the kk-th and (k+1)(k+1)-th expert scores, divided by the spread of that token's scores. A near-tie is a margin below 0.01.

PairAUC at step 8Flips at step 16AUC at step 16Median margin (flipped)Median margin (kept)
Layout DP2 (fp32)1.0000.08%0.9700.00020.073
Layout DP8/EP4 (fp32)1.0000.06%0.9850.00010.073
1-ULP nudge (fp32)1.0000.06%0.9680.00010.073
Per-rank capacity (fp32)0.9723.36%0.9460.00460.077
Per-rank balancing loss (fp32)0.9702.52%0.9500.00370.076

Key insight: Flipped decisions have margins 380–700× smaller than unchanged ones, and the 10% of decisions with smallest margins flip 65–130× more often than the rest.

Finding 4: More Experts → Less Reproducible Routing, Same Output Drift

ExpertsNear-tie choicesKL at endvs. own seedRouting differs
8 (top-2, width 512)2.7%0.2740.8068%
16 (top-4, width 256)6.6%0.2840.8194%
32 (top-8, width 128)11.9%0.2920.85100%
64 (top-16, width 64)23.3%0.2720.81100%
Dense (width 1024)-0.2060.58-

Near-ties approximately double each time expert count doubles, but output drift stays at ~0.8× the seed level.

Finding 5: MoE vs. Dense, and the Role of Precision

  • Dense in fp32: barely moves (0.12× seed level; zero loss barrier — same solution)
  • Dense in fp16: moves more (0.58× seed level)
  • MoE: reaches ~0.8× at both precisions, reaching the 0.01 difference level at step 64 (vs. step 128 for dense)

"The router makes MoE divergence close to all-or-nothing."

Finding 6: Layout Switch Timing Matters

Switch pointFinal KLvs. seedRouting differs
10% of pretraining (step 128)0.12–0.13~35%58–61%
84% of pretraining (step 1,024)0.007~2%15%

Both cases show zero loss barrier, meaning the resumed model stays in the same valley. A layout switch behaves exactly like a 1-ULP nudge at both points.

Model Size Study

At three sizes (25M, 68M, 117M params, fp16), the layout difference is 0.81, 0.88, and 0.86 of seed level. Near-tie share grows slightly with size (6.6%, 6.8%, 7.8%), and larger models diverge earlier.


Theoretical and Practical Implications

Theoretical Implications

  1. MoE divergence is precision-independent: Unlike dense models, where the magnitude of rounding noise determines drift, MoE models diverge to ~80% of seed level at any precision because even the smallest perturbation can flip a near-tie decision.

  2. Near-ties are the amplification mechanism: The router acts as a sensitive amplifier that converts tiny arithmetic differences into qualitatively different models.

  3. Fine-grained experts are functionally interchangeable: More experts create more near-ties and less reproducible routing, but output drift doesn't increase—suggesting individual expert identity matters less than the ensemble.

  4. Sensitivity is phase-dependent: Training is chaotic early on (matches Kwok et al., 2025) but becomes robust to perturbations in the late phase.

Practical Implications

  1. Reproducibility requires fixed layouts: If you need the same model back (debugging, auditing, fair comparisons), keep the layout fixed and use bitwise-deterministic training.

  2. Compare methods on the same layout with several seeds: A layout change moves an MoE almost as much as a seed change—differences smaller than seed-to-seed spread could be pure layout effects.

  3. Schedule layout switches late: For elastic training, switch after the loss flattens rather than right after warm-up.

  4. Monitor the router, not just the loss: Near-tie share and routing decision changes are cheap to log and reveal reproducibility issues that loss alone cannot show.

  5. Plan for routing fragility: Applications depending on specific experts (RL routing replay, expert pruning, interpretability) need deterministic fixed-layout training or a recorded routing table.


Conclusion

The paper demonstrates that the parallelism layout does not change MoE model quality but does change which model you get—about 80% as much as a new random seed, with most routing decisions different. The cause is the router: near-tie decisions flip under the smallest rounding differences, and runs then drift apart on their own. Fine-grained experts create more near-ties and make routing less reproducible without worsening output drift. A dense model of the same active size is far more stable and depends on numerical noise magnitude, while the MoE diverges at any precision.

Future directions implied by this work include:

  • Testing at production scale and with bf16/FP8 precision
  • Investigating layout switches during learning-rate decay, mid-training, or post-training
  • Studying tensor and pipeline parallelism effects
  • Developing deterministic training methods that preserve routing decisions across layouts

"For anyone who needs the same MoE twice, the layout belongs in the experiment record, and fixed-layout deterministic training is the only reliable way to get it."

Related papers