Ignoring systems efficiency leads MoE scaling laws to push sparsity to a boundary, but a systems-aware co-design (MOSAIC) finds interior, cluster-optimal sparsities that lower training loss and better use real cluster compute, a result validated on NVIDIA B200 clusters and staged pretraining up to 250B parameters.

Main Finding

Compute-optimal (model-FLOPs-only) scaling prescriptions for sparse Mixture-of-Experts (MoE) models push sparsity to the data-fit boundary (i.e., “sparser is always better” within the calibrated range). That boundary solution often cannot be executed within real cluster/time budgets because architecture affects attainable throughput. MOSAIC — a joint scaling-law + calibrated performance model + discrete-geometry optimizer — finds cluster-aware (“deliverable-FLOPs”) optima that can be interior in sparsity and reduce training loss under a fixed hardware/time envelope. In short: compute-optimal is not cluster-optimal; architecture and systems co-design matters for realistic training cost and final loss.

Key Points

• Problem: Conventional scaling laws optimize loss given a model-FLOPs budget Cmodel = 6 · Nact · D, ignoring that hardware delivers only a fraction of peak FLOPs and that this fraction depends on architecture and parallel layout.
• MFU (Model FLOPs Utilization): fraction of cluster peak throughput actually delivered to the model’s mathematical FLOPs. MFU captures kernel quality, communication/overhead, pipeline bubbles, dispatch/combines, and expert-parallel all-to-all effects.
• Deliverable budget: Cdeliverable(Z,P, B) = Cpeak · MFU(Z,P,B) · ηgood(Ndev), where Cpeak = Fpeak · Ndev · Ttrain and ηgood is goodput (time spent actually training). Feasibility requires Cmodel ≤ max_{feasible plans} Cdeliverable.
• New scaling law: a four-dimensional MoE scaling law was fit that explicitly includes sparsity S (fraction of parameters inactive per token), plus other MoE geometry dims. Within the calibrated sparsity range, loss decreases monotonically with increasing sparsity — i.e., the model-FLOPs optimum lies at the supported sparsity boundary.
• Systems flip the optimum: when MFU and memory/communication constraints are priced in, an interior sparsity becomes optimal for a given cluster & training window. MOSAIC jointly optimizes discrete MoE geometries, token budget, and execution layout (tensor/pipe/expert/data/context parallelism) as a mixed-integer nonlinear program (MINLP).
• Tools and validation:
  • An operator-level analytical performance model predicts per-block forward/backward cost and MFU; validated on runs up to 18B active parameters with mean absolute percentage error < 15% per sweep.
  • MOSAIC was applied on NVIDIA B200 (p6-B200) clusters, with staged pretraining validations up to 250B total parameters. Example: under 4 p6-B200 nodes for 5 days, the model-FLOPs-optimal sparsity was infeasible; MOSAIC found a cluster-feasible interior sparsity with ~0.031 nats lower predicted loss than the infeasible boundary choice.
• Practical formulation: MOSAIC enumerates realizable discrete geometries Z ∈ Zdisc, computes feasible execution plans Pfeas(Z,B) under device memory caps, predicts MFU and Cdeliverable for each plan, and selects architecture and token budget minimizing predicted loss subject to Cmodel ≤ max_Pfeas Cdeliverable.

Data & Methods

• Architecture & training recipe:
  • Decoder-only transformer language models; every FFN replaced by an MoE layer (first block dense).
  • Each MoE layer includes one shared expert always active plus K routed experts (top-K routing); routing includes Gumbel routing and load-balancing bias controls.
  • Training: next-token prediction, bf16, AdamW, single-pass over data; smoothed final-checkpoint training loss used for scaling-law fits.
• Scaling-law fit:
  • New joint MoE scaling law with variables exposing systems knobs, explicitly including sparsity S, active parameters Nact, total parameters Ntot, and expert geometry factors (G, E, K as appropriate).
  • Data sweep: active parameters from ~104M up to 2.7B, total model sizes up to 79B. Within this calibrated range, fitted loss decreases monotonically with sparsity.
• Performance model:
  • Analytical, operator-level per-block cost model capturing forward/backward FLOPs, memory footprints, and communication patterns (all-to-all, collectives).
  • Predicts MFU for a given architecture Z and execution plan Pℓ and batch envelope B.
  • Validated against training grids (up to 18B active params): MFU predictions have MAPE < 15% per sweep.
• Optimization:
  • MOSAIC is formulated as a bilevel MINLP: minimize L(Z,D) over discrete architectures and token counts subject to architecture-dependent deliverable-FLOPs constraint (inner maximization over feasible execution plans).
  • Solved by structured enumeration leveraging properties (e.g., closed-form elimination of D) and a realizable architecture geometry ladder.
• Experimental deployments:
  • Applied on NVIDIA B200-based clusters (p6-B200 nodes), example constraints: 4 nodes × 5 days used to illustrate infeasibility of boundary sparsity.
  • Ranked MOSAIC-selected architectures were validated with staged pretraining runs up to 250B total parameters on the authors’ infrastructure.

Implications for AI Economics

• Cost forecasting and investment decisions must account for deliverable (architecture-dependent) FLOPs, not model-FLOPs alone. Using Cmodel to price training (or to trade off N vs D) risks systematic misallocation and infeasible plans.
• Cloud/GPU-hour pricing and benchmarking: two architectures with identical theoretical FLOPs can deliver very different wall-clock progress (loss vs time) because of MFU differences. Pricing models and procurement decisions should incorporate MFU-based goodput estimates.
• Design incentives and R&D allocation:
  • MoE and other sparsity techniques may reduce per-token activated FLOPs but raise memory, communication, and synchronization costs — changing marginal returns to parameter investments.
  • Without systems-aware accounting, scaling laws will bias R&D toward architectures that look compute-cheap but are hard/expensive to realize on existing clusters.
• Resource allocation and market dynamics:
  • Frontier labs and cloud providers should jointly optimize hardware procurement and model architectures — e.g., buying network bandwidth or different device types matters as much as raw TFLOPs for MoE efficiency.
  • MFU and deliverable-FLOPs become first-order quantities for capacity planning, time-to-train trade-offs, and competitive dynamics (who can actually realize theoretical gains fastest).
• Policy and forecasting implications:
  • Economic and policy models that rely on compute budgets (e.g., forecasting AI progress, measuring compute concentration) should incorporate architecture-dependent utilization factors; otherwise predictions of achievable scale and timelines may be overly optimistic.
• Carbon and cost efficiency:
  • Architecture choices that minimize model-FLOPs may not minimize wall-clock energy or cost if MFU collapses; therefore, environmental and financial footprint assessments must use deliverable-FLOPs or time-to-loss metrics.
• Practical advice for practitioners and economists:
  • Use a performance-model + scaling-law co-design loop (MOSAIC-like) before committing large training budgets.
  • Report MFU and goodput alongside theoretical FLOPs in benchmarks and economic analyses.
  • Consider cluster-tailored geometry choices (hidden sizes, expert counts, parallel layouts) as adjustable investments that affect both price and produced model quality per dollar.

Caveats and limits - Results are calibrated on a specific MoE recipe, routing/load-balancing choices, training stack (Megatron-Core style), and hardware family (B200/p6-B200). Generalization to other hardware or substantially different MoE implementations requires re-fitting the scaling law and re-validating the performance model. - The scaling-law sparsity behavior is reported within the calibrated data support (active params 104M–2.7B, Ntot up to 79B); extrapolation beyond that range is uncertain. - Work focuses on training-time co-design; inference-time costs and economics are left for future work.