Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts

Summary (Overview)

• MOSAIC framework: The paper introduces MOSAIC (Model Optimization via Systems-Aware TraIning Co-design), a unified optimization framework that jointly selects model architecture, training-token budget, and distributed parallel execution layout for sparse Mixture-of-Experts (MoE) language models, rather than treating scaling laws and systems tuning as disconnected stages.

• Key finding on sparsity: Under a model-FLOPs budget alone, MoE sparsity admits no interior optimum—the fitted loss decreases monotonically with sparser models, driving the optimum to the upper boundary of the calibrated sparsity range. An interior, meaningful optimal sparsity emerges only when hardware-deliverable FLOPs (accounting for MFU, communication, and memory constraints) replace raw model FLOPs as the budget.

• New scaling law: The authors fit a four-dimensional joint scaling law $\mathcal{L}(N_{\mathrm{tot}}, S, D, G)$ over total parameters, sparsity, training tokens, and expert split factor, spanning active parameters from 104M to 2.7B and total model sizes up to 79B parameters. The expert split factor $G$ enters with an identified exponent $\eta \approx 0.95$.

• Performance model validation: An operator-level analytical performance model (released as ScalePlan) predicts MFU with mean absolute percentage error under 15% across sweeps up to 18B active parameters, with 100% Top-5 overlap in configuration ranking for large models.

• Empirical validation: Staged pretraining runs up to 250B total parameters confirm that the loss ordering flips between the model-FLOPs axis and the hardware-compute axis, validating the co-design principle that compute-optimal configurations are not necessarily cluster-optimal.

Introduction and Theoretical Foundation

Background and Motivation

Traditional large-scale pretraining decouples architecture selection from systems optimization:

1. Scaling law stage: Selects architecture and training recipe by optimizing loss under a compute constraint (e.g., Chinchilla-style laws over model size $N$ and token count $D$).
2. Systems stage: Optimizes implementation for hardware efficiency, taking model size as given.

This separation is consequential because two architectures with the same predicted quality can differ several-fold in training throughput once attainable MFU and best parallel layout are accounted for. The paper argues that the architecture chosen by a scaling law should account for systems efficiencies rather than model FLOPs alone.

Theoretical Foundation

The central theoretical contribution is the concept of deliverable model FLOPs, which replaces the architecture-independent model-FLOPs budget:

$$C_{\mathrm{deliverable}}(Z, P_\ell, B) = C_{\mathrm{peak}} \cdot \mathrm{MFU}(Z, P_\ell, B) \cdot \eta_{\mathrm{good}}(N_{\mathrm{dev}})$$

where $C_{\mathrm{peak}} = F_{\mathrm{peak}} N_{\mathrm{dev}} T_{\mathrm{train}}$ is the raw hardware ceiling, MFU is the Model FLOPs Utilization, and $\eta_{\mathrm{good}}$ is the goodput fraction.

MFU is defined as:

$$\mathrm{MFU}(Z, P_\ell, B) = \frac{6 N(Z) B_{\mathrm{global}} T_{\mathrm{seq}}}{t_{\mathrm{iter}}(Z, P_\ell, B) N_{\mathrm{dev}} F_{\mathrm{peak}}}$$

The key insight is that $C_{\mathrm{deliverable}}$ is architecture-dependent: different architectures achieve different MFU values, so a fixed hardware budget translates into different amounts of useful computation depending on the chosen model geometry.

Key Definitions

The paper introduces two architectural knobs that parameterize the sparse MoE scaling laws:

• Sparsity: $S = 1 - N_{\mathrm{act}}/N_{\mathrm{tot}}$, the fraction of parameters inactive per token.
• Expert split factor: $G := d_{\mathrm{ff}}/d_{\mathrm{expert}}$, the ratio of the dense FFN width to a single expert's hidden width, measuring how finely the dense FFN is partitioned into experts.

Methodology

MOSAIC Optimization Formulation

MOSAIC solves a bilevel mixed-integer nonlinear program (MINLP):

$$\min_{Z \in \mathcal{Z}_{\mathrm{disc}}, D} \mathcal{L}(Z, D) \quad \text{subject to} \quad C_{\mathrm{model}}(Z, D) \leq \max_{P_\ell \in \mathcal{P}_{\mathrm{feas}}(Z, B)} C_{\mathrm{deliverable}}(Z, P_\ell, B)$$

where:

• $Z$ is the architecture decision variable (a realizable sparse MoE geometry)
• $D$ is the number of training tokens
• $P_\ell = (P_{\mathrm{TP}}, P_{\mathrm{EP}}, P_{\mathrm{CP}}, P_{\mathrm{PP}}, P_{\mathrm{DP}}, B_{\mathrm{micro}}, A_{\mathrm{ckpt}})$ is the execution plan
• $\mathcal{P}_{\mathrm{feas}}(Z, B)$ collects execution plans satisfying memory constraints: $\mathrm{Mem}(Z, P_\ell, B) \leq M_{\mathrm{dev}}$

Joint Scaling Law

The fitted scaling law incorporates total parameters, sparsity, training tokens, and expert split factor:

$$\mathcal{L}(N_{\mathrm{tot}}, S, D, G) = \frac{a}{N_{\mathrm{tot}}^{\alpha}} + \frac{b}{D^{\beta}} + \frac{c}{(1-S)^{\lambda}} + \frac{j}{(1-S)^{\delta} N_{\mathrm{tot}}^{\gamma} G^{\eta}} + e$$

The fitted coefficients (with 95% BCa confidence intervals) are:

The expert split factor exponent $\eta \approx 0.95$ is a robust observation: larger $G$ discounts the loss penalty that sparsity imposes at fixed capacity.

Performance Model (ScalePlan)

The performance model estimates MFU and memory using:

1. Microbenchmarks on target hardware for: dense GEMMs, grouped GEMMs (MoE expert MLP), SDPA attention, top-k routing, cross-entropy, NCCL collectives, and CPU dispatch gaps.

2. Iteration-time model with three key components:

   • Compute vs. dispatch: $t_{\mathrm{mb}} = \max(t_{\mathrm{cmp}}, t_{\mathrm{disp}})$ for the per-microbatch wall time, where $t_{\mathrm{disp}} = L_{\mathrm{stage}}\kappa(E_{\mathrm{loc}})t_{\mathrm{gap}} + t_{\mathrm{host}}$
   • Pipeline bubble: $f_{\mathrm{bubble}} = \frac{P_{\mathrm{PP}} - 1}{P_{\mathrm{VPP}} n_{\mathrm{mb}}}$ under 1F1B scheduling
   • Vocabulary-stage imbalance: $t_{\mathrm{vocab}} = (n_{\mathrm{mb}} - 1)\max(t_{\mathrm{first}}, t_{\mathrm{last}}) + t_{\mathrm{first}} + t_{\mathrm{last}}$

3. Single calibration term capturing cross-node straggler and overlap effects:

$$\chi_{\mathrm{sys}} = 1 + c_{\mathrm{A2A}} \frac{t_{\mathrm{A2A}}}{t_{\mathrm{iter}}} \mathbb{1}[N_{\mathrm{node}}^{\mathrm{EP}} > 1] + c_{\mathrm{PP}}(N_{\mathrm{node}}^{\mathrm{PP}} - 1) + c_{\mathrm{ovl}} \mathbb{1}[P_{\mathrm{PP}} \geq N_{\mathrm{gpu/node}}, N_{\mathrm{node}}^{\mathrm{PP}} = 1]$$

Geometry Scaling Ladder

The search space is structured via a geometry ladder where layer count and head count follow $L_{\mathrm{layers}} = n_{\mathrm{head}} = mq$ with multiplier $m = 4$, residual width $d = n_{\mathrm{head}} \cdot 128$, and reference dense FFN width $d_{\mathrm{ff}} = 2.5d$. A single seed $q$ fixes the backbone geometry.

Empirical Validation / Results

Scaling Law Findings

Boundary-seeking sparsity: Under a fixed model-FLOPs budget $C_{\mathrm{model}} = 6 N_{\mathrm{act}} D$, the optimal sparsity lies at the upper boundary $S = S_{\mathrm{max}}$ of the calibrated range. Three monotone trends hold over $S \in [0.5, 0.981]$:

• Optimal total parameters $N_{\mathrm{tot}}^{\star}(S)$ increases monotonically with $S$
• Optimal active parameters $N_{\mathrm{act}}^{\star}(S)$ decreases
• Optimal loss $\mathcal{L}^{\star}(S)$ decreases monotonically at every $C_{\mathrm{model}}$

The joint law achieves held-out RMSE of 0.00779 in loss units on the largest 10% of runs by $N_{\mathrm{act}}$ within each $G$ stratum, with $R^2 = 0.99877$.

Performance Model Validation

Validation spans three sweeps on AWS p6-b200.48xlarge nodes (8 NVIDIA B200 GPUs each):

Top-5 overlap between predicted and measured MFU rankings is 100% for 5.4B and 18B sweeps, 80% for 700M. The dominant systems lever is the pipeline fill ratio $r_{\mathrm{fill}} = n_{\mathrm{mb}}/P_{\mathrm{PP}}$, explaining up to 3.1× layout spread in MFU.

MOSAIC Optimization Results

Optimal configuration under fixed envelope (32 p6 nodes, 20 days, $C_{\mathrm{peak}} = 9.95 \times 10^{23}$ FLOPs):

Key insight: The candidate that realizes the most model FLOPs ($1.94 \times 10^{23}$ FLOPs, a G=2, E=32 geometry) is NOT the lowest-loss configuration. The optimal configuration reaches $\mathcal{L} = 1.3898$ on $1.23 \times 10^{23}$ model FLOPs—roughly 36% fewer than the rightmost candidate. Expert split factor $G$ orders the spread: G=8 geometries sit at low $C_{\mathrm{model}}$ (median MFU ~8%), while G=2 geometries reach the right end (13–20% MFU).

Budget Sweeps

• Optimal sparsity $S^{\star}$ sits on a wide plateau (roughly 0.915–0.963) and does not sharpen with budget—no monotone sparsity trend under hardware constraints (unlike model-FLOPs optimality).
• Optimal expert split factor $G^{\star}$ is 4 across most budgets, with G=8 appearing only at the longest window and G=2 at the largest node budget.
• Active parameters stay near 14–15B across the days range, growing to 34B and 58B at the two largest node budgets.

Staged Pretraining Validation

Four configurations were selected from MOSAIC's ranked predictions on an 8-node, 10-day envelope:

Results:

1. The predicted MFU ranking reproduces exactly on hardware (rank correlation aligned).
2. On the model-FLOPs axis, ID 2 (G=8, low active count) holds an advantage over much of the range; on the hardware-compute axis ($C_{\mathrm{cost}} = C_{\mathrm{model}}/\mathrm{MFU}$), ID 1 (G=4, higher MFU) leads—the loss ordering flips between axes, confirming the co-design effect.

Theoretical and Practical Implications

Theoretical Contributions

1. Reframing the budget: The paper demonstrates that model FLOPs is an inadequate fairness metric for sparse MoE architecture selection. The appropriate budget is hardware-deliverable FLOPs, which is architecture-dependent through the MFU term.

2. Emergence of interior optima: The boundary-seeking sparsity under model-FLOPs budgets becomes an interior optimum once systems constraints (memory floor, GPU-hour costs, all-to-all penalties) are imposed. This emerges naturally from constraints rather than explicit caps on model dimensions.

3. Identified scaling exponent: The expert split factor exponent $\eta \approx 0.95$ is identified (not on the coefficient ridge), quantifying how finer-grained experts soften the sparsity penalty at fixed memory footprint—a benefit that must be priced against communication overhead.

Practical Implications

1. Unified co-design: Practitioners should select architectures and parallel layouts jointly rather than in disconnected stages. The Kimi-K2 approach of treating MoE scaling laws and systems tuning separately is explicitly called out as consequential at scale.

2. Design rules for MoE parallelism: The MFU sweeps yield concrete rules: the pipeline fill ratio $r_{\mathrm{fill}} = n_{\mathrm{mb}}/P_{\mathrm{PP}}$ is the dominant lever; expert parallelism helps up to the node boundary and hurts past it; fix $P_{\mathrm{TP}}=1$, set $P_{\mathrm{EP}}$ to the scale-up domain, and keep $r_{\mathrm{fill}} \geq 4$.

3. Hardware-dependent optimal geometry: The optimal sparsity and expert split factor are functions of cluster size and training window, not universal constants. Larger hardware budgets support larger active parameter counts.

Conclusion

Main Takeaways

MOSAIC demonstrates that compute-optimal configurations are not necessarily cluster-optimal for sparse MoE language models. By coupling a four-dimensional scaling law with an operator-level performance model, the framework shows that:

1. Model-FLOPs-only optimization drives sparsity to the boundary of the calibrated range
2. Hardware constraints (memory, communication, MFU) create an interior optimal sparsity
3. The architecture that minimizes loss under a fixed hardware envelope is not the one that emits the most model FLOPs
4. Loss ordering between configurations flips between model-FLOPs and hardware-compute axes, validating the co-design principle

Limitations

• Results are conditional on the chosen geometry ladder, hardware platform (NVIDIA B200), and training recipe
• Some reported optima require substantial extrapolation beyond the scaling law training regime
• Uncertainty in non-identified scaling law coefficients is not propagated through the discrete search
• The performance model mirrors a basic Megatron-Core stack; the absence of monotone trend in optimal $G$ may be an artifact of that stack's systems penalty
• Expert split factor is limited to $G \leq 8$
• The overtrained regime (beyond tokens-per-parameter ratio ~70) is not studied

Future Directions

1. Joint batch-size and parallelism optimization: Extend the decision space to include global batch size
2. Broader geometry families: Test whether the $(N_{\mathrm{tot}}, S, G)$ representation transfers across backbone geometries, incorporating hardware-friendly tile/wave constraints
3. Communication-computation overlap: Extend the performance model to represent persistent/fused kernels that overlap collectives with computation
4. Test-time scaling laws: Apply the co-design framework to inference under hardware costs