🤖 machine learning

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

This paper introduces MOSAIC, a systems-aware co-design framework that demonstrates how optimal sparsity for sparse Mixture-of-Experts models emerges only when jointly optimizing model architecture with hardware constraints, rather than through traditional compute-optimal scaling laws alone.

Original authors: Soumajyoti Sarkar, Yuxin Tang, Sheng Zha

Published 2026-08-12

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Soumajyoti Sarkar, Yuxin Tang, Sheng Zha

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build the ultimate, most powerful robot brain. In the world of artificial intelligence, this "brain" is a massive computer program called a Large Language Model (LLM). To make these brains smarter, scientists have discovered a simple rule: if you feed them more data and make them bigger, they get better at understanding the world. This rule is called a "scaling law." For a long time, the recipe for building these brains was split into two separate steps. First, a mathematician would figure out the perfect size for the brain based on how much computer power (called "compute") was available. Then, a different engineer would try to squeeze that brain onto a specific set of computer chips to make it run as fast as possible.

Think of it like planning a road trip. The first step is deciding how far you want to drive based on your gas budget. The second step is picking the car. But what if the car you picked is so heavy and clunky that it sips gas at a terrible rate, meaning you can't actually drive as far as you planned? That is the problem this paper tackles. It argues that you can't just pick a brain size and then worry about the car. You have to design the brain and the car together, because the shape of the brain changes how efficiently the car drives. If you ignore the car's engine, you might end up with a "perfect" brain that is too heavy to move, leaving you stuck at the starting line.

The Problem: The "Perfect" Brain That Won't Fit

The paper introduces a new way of thinking called MOSAIC (Model Optimization via Systems-Aware TraIning Co-design). The authors, working at Amazon AGI Foundations, noticed that the old way of doing things was leaving money on the table. Traditionally, researchers would calculate the "Compute-Optimal" model. This is the model size that gets the best results for a fixed amount of computer power, assuming every bit of power is used perfectly.

However, the authors found a hidden trap. They focused on a special type of AI architecture called a Mixture-of-Experts (MoE). Imagine a giant library where, instead of one giant librarian reading every book, you have thousands of tiny, specialized experts. When a question comes in, a smart router sends it to just a few of these experts (say, 2 out of 100) to answer. This makes the model very fast and efficient because it doesn't have to wake up the whole library for every question.

The old math suggested that to get the best results, you should make the library as huge as possible and wake up as few experts as possible. In other words, the "perfect" model should be incredibly sparse (mostly empty). The paper shows that if you only look at the raw math of "how many calculations do we need?", the answer is always: "Make it as sparse as you can." The math says the best sparsity is at the very edge of what is possible.

The Twist: The Car Engine Matters

Here is where the story changes. The authors realized that while the math says "make it super sparse," the physical reality of the computer chips says "no way."

When you have a model with thousands of tiny experts, the computer has to do a lot of extra work just to decide which experts to wake up and to pass the information between them. It's like having a giant office building where the manager spends all day running between cubicles to deliver memos, leaving very little time for actual work. The more sparse the model gets, the more time the computer wastes on these "running around" tasks (called communication costs) and the less time it spends on the actual thinking (calculations).

The paper argues that the old "Compute-Optimal" math was flawed because it assumed the computer could deliver 100% of its power to the thinking part. In reality, a super-sparse model might only get 8% of the computer's power to do actual thinking, while a slightly denser model might get 15%.

The Solution: MOSAIC

To fix this, the team built MOSAIC. Instead of asking, "What is the best model for a fixed amount of math?", they asked, "What is the best model we can actually run on our specific cluster of computers in a fixed amount of time?"

They combined two things:

A Scaling Law: A prediction of how smart the model will be based on its size and shape.
A Performance Model: A simulator that predicts how fast the model will actually run on real hardware, accounting for the "running around" costs.

When they ran MOSAIC, the results were surprising. The "perfect" model wasn't the one at the extreme edge of sparsity. Instead, the best model was an interior solution—a happy medium. It was sparse enough to be efficient, but not so sparse that the computer got bogged down by the overhead of managing all those tiny experts.

What They Found

The team tested this on real hardware using NVIDIA B200 GPUs. They trained models ranging from 700 million to 18 billion active parameters (and up to 79 billion total parameters).

The Old Way: If you just followed the old math, you would pick a model with a sparsity of about 0.985 (meaning 98.5% of the experts are asleep). But on their specific cluster, this model would be so slow that it would take forever to train, and it would actually end up less smart than a slightly denser model.
The MOSAIC Way: The system found that the sweet spot was a sparsity of around 0.96. This model wasn't the "theoretical" best, but it was the practical best. It ran faster, used the computer's power more efficiently, and achieved a lower error rate (meaning it was smarter) than the "theoretical" best.

In fact, they showed that the model that looked best on paper (the one with the most "Model FLOPs") was actually the slowest to train in real life. The model that won was the one that balanced the math with the reality of the hardware.

Why This Matters

This paper suggests that we need to stop treating the design of AI models and the design of the computers that run them as separate steps. You can't just pick a model size and hope the computers can handle it. The shape of the model changes how the computer behaves.

The authors are careful to note that their findings are specific to the type of computer chips they used (NVIDIA B200s) and the specific software they built. They didn't find a universal law that applies to every computer in the universe. However, they proved that for the frontier of AI training, ignoring the "systems" (the hardware and how it moves data) leads to suboptimal results.

By using MOSAIC, they showed that the "best" model is not a fixed point on a graph, but a moving target that depends on your specific hardware budget. The most efficient model is the one that fits your cluster like a glove, not the one that looks best on a piece of paper. This is a shift from "Compute-Optimal" (what the math says) to "Cluster-Optimal" (what the hardware can actually deliver).

In short, if you want to build the smartest AI possible, you don't just need a bigger brain; you need a brain that fits the body it lives in.

Technical Summary: Compute-Optimal Is Not Cluster-Optimal

Problem Statement
Current large-scale pretraining workflows conventionally separate algorithmic scaling decisions from systems implementation. A "scaling law" stage selects a model architecture and data budget to minimize loss under a fixed compute (FLOPs) constraint, followed by a disconnected "systems" stage that optimizes the implementation for hardware efficiency. The authors argue this separation is flawed for sparse Mixture-of-Experts (MoE) models. While traditional scaling laws optimize for Model FLOPs ( $C_{model}$ ), they ignore the systems efficiency of the prescribed architecture. Two architectures with identical predicted loss and $C_{model}$ can yield vastly different training throughputs due to variations in Model FLOPs Utilization (MFU), communication costs, and memory footprints. Consequently, the "compute-optimal" configuration derived from a FLOPs budget may be infeasible or suboptimal on a specific hardware cluster, as it fails to account for the "deliverable" compute a cluster can actually sustain over a fixed training window.

Methodology: MOSAIC
The authors introduce MOSAIC (Model Optimization via Systems-Aware Training Co-design), a framework that formulates model architecture and systems co-design as a single optimization problem. MOSAIC jointly selects the model architecture, training token budget, and distributed execution layout (parallelism strategy) to minimize predicted loss under a hardware-constrained budget.

The framework integrates three core components:

A Four-Dimensional Joint Scaling Law: The authors fit a new scaling law for sparse MoE models that includes the total parameter count ( $N_{tot}$ ), training tokens ( $D$ ), sparsity ( $S$ ), and an expert split factor ( $G$ ). $G$ is defined as the ratio of the dense feed-forward width to the single expert width ( $G = d_{ff}/d_{expert}$ ), capturing the granularity of the experts. The law is formulated as:
$L(N_{tot}, S, D, G) = \frac{a}{N_{tot}^\alpha} + \frac{b}{D^\beta} + \frac{c}{(1-S)^\lambda} + \frac{j}{(1-S)^\delta N_{tot}^\gamma G^\eta} + e$
This formulation allows the loss to be predicted based on discrete geometric choices rather than abstract parameter counts.
An Operator-Level Performance Model: To estimate systems efficiency, the authors developed an analytical performance model that predicts MFU and memory footprint for a given architecture and parallelism layout. This model accounts for:
- Compute costs (dense and grouped GEMMs, attention).
- Communication overheads (all-to-all for expert parallelism, pipeline bubbles, data-parallel collectives).
- CPU dispatch latency (critical for small micro-batches).
- Hardware-specific calibration factors for inter-node stragglers and intra-node overlap.
  The model is validated against training runs up to 18 billion active parameters, achieving a mean absolute percentage error (MAPE) of under 15% in MFU prediction.
Joint Optimization (MINLP): MOSAIC solves a mixed-integer nonlinear program (MINLP) that minimizes the predicted loss subject to the constraint that the required Model FLOPs ( $C_{model}$ ) do not exceed the Deliverable Model FLOPs ( $C_{deliverable}$ ).
$C_{deliverable} = C_{peak} \times \text{MFU}(Z, P_\ell, B) \times \eta_{good}$
Here, $C_{peak}$ is the raw hardware ceiling, and MFU is architecture-dependent. The optimization searches a discrete grid of realizable geometries ( $Z$ ) and execution plans ( $P_\ell$ ) to find the configuration that minimizes loss within the cluster's actual capabilities.

Key Contributions

Reformulation of the Compute Budget: The paper demonstrates that for sparse MoE models, optimizing under a fixed $C_{model}$ budget leads to a "boundary-seeking" behavior where the optimal sparsity lies at the upper limit of the data support (i.e., the sparsest possible model). This is because the scaling law predicts lower loss for sparser models at fixed FLOPs, ignoring the systems cost.
Systems-Aware Co-Design: By coupling the scaling law with a performance model, MOSAIC shifts the budget from $C_{model}$ to $C_{deliverable}$ . This reveals that the optimal sparsity is an interior point determined by the cluster's specific constraints (memory, bandwidth, parallelism overheads), rather than the boundary.
Performance Model Validation: The authors provide a validated, operator-level performance model that accurately predicts MFU and identifies the optimal parallel layout (Tensor, Expert, Pipeline, Data Parallelism) for sparse MoEs, showing that the "best" layout is highly dependent on the specific geometry and cluster size.
Empirical Validation: The framework was tested on AWS p6 nodes (NVIDIA B200 GPUs) with runs up to 250B total parameters. The results confirmed that the configuration ranking based on Model FLOPs differs significantly from the ranking based on hardware-compute efficiency (wall-clock time).

Results

Boundary vs. Interior Optima: Under a fixed $C_{model}$ , the fitted loss decreases monotonically with increasing sparsity, pushing the optimum to the boundary ( $S_{max}$ ). However, under a fixed hardware budget ( $C_{deliverable}$ ), the loss curve exhibits an interior minimum. The systems costs (communication, memory) of extremely sparse models eventually outweigh their theoretical loss benefits.
Loss vs. Throughput Trade-off: In experiments with a 32-node, 20-day budget, the configuration that realized the most Model FLOPs (a denser, coarser expert geometry) did not achieve the lowest loss. Conversely, the loss-optimal configuration (a sparser, finer expert geometry) sustained a lower MFU but achieved a lower final loss because it was the only configuration that fit within the deliverable FLOPs constraint while maintaining a favorable loss trajectory.
Real-World Validation: Staged pretraining runs on 8-node clusters validated the framework. The measured MFU ordering matched the predicted ordering. Furthermore, when plotting loss against "peak-equivalent hardware compute" (accounting for MFU), the ranking of configurations flipped compared to the ranking against raw Model FLOPs, confirming that the "compute-optimal" model is not the "cluster-optimal" model.

Significance
The paper argues for a fundamental shift in how frontier language models are trained: moving from disconnected stages of "scaling law selection" and "systems tuning" to a unified architecture-systems co-design. The authors claim that relying solely on Model FLOPs budgets is insufficient for sparse MoE models because it fails to price in the systems efficiency of the architecture. By introducing MOSAIC, they demonstrate that the optimal model geometry is not a universal constant but is contingent on the specific hardware cluster, the training window, and the parallelism constraints. This approach ensures that the selected architecture is not only theoretically loss-optimal but also practically feasible and efficient on the target hardware.

Technical Summary: Compute-Optimal Is Not Cluster-Optimal

More like this