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:
- Numerical effects: Computers add numbers with finite precision, so , and changing layout changes the order of gradient and expert-output additions.
- 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
- Does the parallelism layout change the trained MoE model?
- If so, why?
- 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
| Item | Setting |
|---|---|
| Base MoE | 6 layers, width 256, 4 heads, RoPE, RMSNorm, QK-Norm, no QKV bias, tied embeddings |
| Experts | 16 SwiGLU experts (width 256), top-4 routing, no shared expert; ~25M total / ~10M active params |
| Dense control | Same model, one SwiGLU block of width 1024 (same active size) |
| Data | FineWeb-Edu; 16,384-token byte-level BPE; 20M training tokens (1,221 steps), global batch 32 × 512 tokens |
| Optimizer | AdamW, peak LR , 100 warm-up steps, weight decay 0.1, , clip 1.0 |
| Precision | fp32 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 < 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
| Comparison | KL at step 8 | KL at end | vs. own seed | Top-1 agree | Routing differs | Loss barrier | Loss change |
|---|---|---|---|---|---|---|---|
| MoE: different seed | 0.350 | 1.00 | 0.664 | - | 3.538 | +0.023 | |
| MoE: different layout (8) | 0.280 | 0.80 | 0.689 | 92% | 0.417 | +0.005 | |
| MoE: 1-ULP nudge (3) | 0.283 | 0.81 | 0.685 | 93% | 0.432 | +0.016 | |
| MoE: thread count | 0.277 | 0.79 | 0.689 | 92% | - | +0.006 | |
| Dense: different layout (2) | 0 | 0.043 | 0.12 | 0.862 | - | - | -0.002 |
| Dense: 1-ULP nudge (2) | 0 | 0.048 | 0.13 | 0.859 | - | 0.000 | -0.002 |
Key observations:
- The difference starts tiny ( 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 ( vs. ), 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 -th and -th expert scores, divided by the spread of that token's scores. A near-tie is a margin below 0.01.
| Pair | AUC at step 8 | Flips at step 16 | AUC at step 16 | Median margin (flipped) | Median margin (kept) |
|---|---|---|---|---|---|
| Layout DP2 (fp32) | 1.000 | 0.08% | 0.970 | 0.0002 | 0.073 |
| Layout DP8/EP4 (fp32) | 1.000 | 0.06% | 0.985 | 0.0001 | 0.073 |
| 1-ULP nudge (fp32) | 1.000 | 0.06% | 0.968 | 0.0001 | 0.073 |
| Per-rank capacity (fp32) | 0.972 | 3.36% | 0.946 | 0.0046 | 0.077 |
| Per-rank balancing loss (fp32) | 0.970 | 2.52% | 0.950 | 0.0037 | 0.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
| Experts | Near-tie choices | KL at end | vs. own seed | Routing differs |
|---|---|---|---|---|
| 8 (top-2, width 512) | 2.7% | 0.274 | 0.80 | 68% |
| 16 (top-4, width 256) | 6.6% | 0.284 | 0.81 | 94% |
| 32 (top-8, width 128) | 11.9% | 0.292 | 0.85 | 100% |
| 64 (top-16, width 64) | 23.3% | 0.272 | 0.81 | 100% |
| Dense (width 1024) | - | 0.206 | 0.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 point | Final KL | vs. seed | Routing 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
-
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.
-
Near-ties are the amplification mechanism: The router acts as a sensitive amplifier that converts tiny arithmetic differences into qualitatively different models.
-
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.
-
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
-
Reproducibility requires fixed layouts: If you need the same model back (debugging, auditing, fair comparisons), keep the layout fixed and use bitwise-deterministic training.
-
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.
-
Schedule layout switches late: For elastic training, switch after the loss flattens rather than right after warm-up.
-
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.
-
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
- SkillDuet: Co-Evolving Models and Skills through Response Feedback
SkillDuet co-evolves model parameters and reusable textual skills via response-based Strength scoring, achieving state-of-the-art 98.57% on ALFWorld and 92.97% on WebShop.
- What Limits Recursive Reasoning Models: Optimization, Architecture and Test-Time Scaling
Recursive reasoning performance is driven primarily by optimization choices—gradient horizon, physical batch size, and stabilized state updates—rather than architectural hierarchy, yielding a 13.6M-parameter model that doubles OOD accuracy on arithmetic.
- Library of Layers (LoL): Adaptively Ordering Transformer Layers at Runtime
Library of Layers lets a router dynamically reorder transformer layers per input, matching dense, looped, and sparse MoE baselines while enabling interpretable, controllable routing paths.