Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts
Soumajyoti Sarkar, Yuxin Tang, Sheng Zha
cs.LG, cs.AI
2026-08-11
MOSAIC co-designs MoE sparsity with hardware. Compute-optimal sparsity sits at a boundary clusters can't reach; the true optimum, 0.031 nats lower, appears under cluster constraints.
Pretraining a sparse MoE (Mixture-of-Experts) model is conventionally a two-stage affair: a scaling-law stage picks the architecture (expert count, sparsity) by minimizing loss under a compute budget, then a systems stage picks the parallelism layout for the hardware. The two stages do not talk.
The catch is that model FLOPs are not the FLOPs the hardware actually delivers. The paper makes a sharp observation: two architectures with identical predicted loss can sustain several-fold different training throughput once you account for the MFU (Model FLOPs Utilization) each can reach and its best parallel layout. Sparse MoE feels this hardest. A scaling law left to itself says sparser is always better, so loss falls monotonically with sparsity and the compute optimum parks at the sparsity ceiling, which on a real cluster you often cannot run at all.
MOSAIC (Model Optimization via Systems-Aware TraIning Co-design), from Amazon's AGI team, puts architecture and systems in the same optimization. For a fixed cluster and training window it jointly picks the model geometry (expert count E, expert split factor G, sparsity S), the token count D, and the full parallel layout P (tensor/expert/context/pipeline/data parallelism, plus micro-batch and activation checkpointing).
The key departure is that the compute budget is not a constant. It is the architecture-dependent Cdeliverable = Cpeak · MFU(arch, layout) · goodput, so the constraint flips from 'Cmodel ≤ a fixed FLOPs figure' to 'Cmodel ≤ Cdeliverable', where Cmodel = 6·Nactive·D is the paper budget and Cdeliverable is what the hardware actually delivers.
Two components feed the optimizer. The first is an extended scaling law L(Ntot, S, D, G) that adds the expert split factor G (G = dff / dexpert, how finely a dense FFN is partitioned). It is fit on 104M to 2.7B active parameters (up to 79B total), sparsity 0.5 to 0.981, with G in {1,2,4,6,8}. The G exponent η≈0.95 is one of only two coefficients the data actually pin down, and its reading is concrete: at fixed memory, finer-grained (larger G) experts reliably soften the loss penalty that sparsity imposes.
The second is a performance model that predicts MFU, memory footprint, and the best parallel layout straight from the model spec and device config. It decomposes an iteration into Transformer compute, pipeline bubble, vocabulary-stage imbalance, collectives, and the optimizer step, with explicit corrections for CPU kernel-launch overhead on small models and cross-node straggler and hop costs on large collectives.
The reason for a single-step discrete search over a 'geometry ladder' instead of continuous-then-round is the boundary problem: the model-FLOPs optimum sits on the boundary, and rounding cannot rescue it.
| Metric | Value |
| Scaling-law fit R² | 0.99877 |
| Held-out RMSE (top 10% per G-stratum) | 0.00779 loss |
| Active params fit | 104M – 2.7B (up to 79B total) |
| MFU error (700M/5.4B/18B sweeps) | 9.6% / 9.0% / 13.2%, all < 15% |
| Layout-rank Top-5 overlap (5.4B/18B vs 700M) | 100% / 80% |
The clearest contrast is Figure 1. At Cmodel = 3.3×10²¹, the pure model-FLOPs optimum is the ceiling Smax = 0.985. On the real hardware of 4 AWS p6-B200 nodes for 5 days, the feasible boundary only reaches S = 0.96; beyond it, configurations deliver fewer FLOPs. MOSAIC finds an interior cluster-optimal sparsity whose predicted loss is 0.031 nats lower than that boundary optimum.
Figure 8 supplies the evidence for the opposite ordering: at 512 devices (64 p6 nodes), the wall-clock per delivered model FLOP rises monotonically with sparsity, up to 1.70× the densest configuration, the reverse of how loss orders with sparsity. That reversal is exactly why the budget must be made architecture-dependent.
For the expert split factor G (Figure 6), predicted loss decreases monotonically with G; observed runs keep decreasing through G=4 and flatten or slightly reverse at G=8, within run-to-run noise.
For anyone training a sparse MoE: choosing the architecture from a scaling law in isolation can hand you a configuration your cluster cannot run efficiently, or one that is worse in real GPU-hours than a denser choice. Architecture and systems are coupled, and co-designing them reaches the same loss in fewer GPU-hours. The framework is practical in that it also returns the parallel layout, not just the model dimensions.
Honestly, this is a co-design and optimization-framework contribution, not a new model or a benchmark SOTA. Its two durable artifacts are the scaling-law extension that adds G and finds η≈0.95 reliably identified, and the conceptual result that no interior sparsity optimum exists under a pure model-FLOPs budget, which is precisely why co-design is needed.