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

# Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts
Source: [https://arxiv.org/html/2608.10605](https://arxiv.org/html/2608.10605)
\\nameYuxin Tang\\emailtanyuxin@amazon\.com \\addrAmazon AGI Foundations \\nameSheng Zha\\emailzhasheng@amazon\.com \\addrAmazon AGI Foundations

###### Abstract

In large\-scale pretraining, the algorithm, architecture, and systems decisions are conventionally made in disconnected stages\. A scaling law stage selects an architecture and training recipe, optimizing loss under compute constraints, and a separate systems stage then optimizes the implementation for hardware efficiency\. In this work, we develop MOSAIC, which formulates model architecture and systems co\-design as an optimization problem\. MOSAIC couples a predictive scaling law with a calibrated performance model that estimates Model FLOPs Utilization \(MFU\), communication cost, memory footprint, and the best parallel layout\. We instantiate the framework for sparse Mixture\-of\-Experts \(MoE\) language models, where expert count, routing sparsity, and other MoE layer dimensions affect both the loss and systems efficiency\. We fit a scaling law on sparse MoE models trained on text data, whose scaling dimensions include the sparsity factor, which is the fraction of model parameters inactive per token in a forward pass\. The scaling law sweeps in our work span active parameters from104104million to2\.72\.7billion and total model sizes reaching7979billion parameters\. We show that, within the calibrated sparsity range, an efficiency\-agnostic model\-FLOPs budget admits no interior optimal sparsity\. The fitted loss decreases monotonically with sparser models and the compute optimum lies at the upper boundary of the data support\. An optimal sparsity in MoE models instead emerges under the cluster’s systems constraints, as captured by MOSAIC\. Our results argue for a shift towards unified architecture and systems co\-design for frontier language model training\.

Keywords:LLM pretraining, scaling laws, hardware\-aware model co\-design

## 1Introduction

Compute\-optimal scaling laws\(kaplan2020scaling;hoffmann2022training\)are the dominant tool for translating a fixed compute budget into a model size and a data count in pretraining\. They are the standard recipe in practice because, once a training compute budget is set, the architecture \(model sizeNN\) and data \(token countDD\) decisions both follow from a single optimization against the predicted\-loss curve, and because the curve extrapolates far past the regime in which it was fit, making it a predictable tool even at frontier scale\(bahri2024explaining;hestness2017deep\)\. For dense models,hoffmann2022trainingdevise a joint scaling law for loss in the model sizeNNand the number of training tokensDD, which admits a closed\-form model FLOPs111We use the term model FLOPs to denote “compute” as is defined in the view of scaling laws and explicitly use the term deliverable model FLOPs to denote the hardware\-delivered FLOPs in training\.optimal allocation\. More recent prescriptive scaling laws recommend the optimal scaling dimensions under additional data constraints\(lovelace2026prescriptive\), yet none of these constraints price in the systems efficiency of the prescribed architecture\. This matters because two architectures of the same predicted quality can realize training throughput that differs several\-fold, once their attainable Model FLOPs Utilization \(MFU\)222MFU is the fraction of a device’s peak throughput that a training run actually sustains\.[https://cloud\.google\.com/blog/products/compute/using\-cloud\-tpu\-multislice\-to\-scale\-ai\-workloads](https://cloud.google.com/blog/products/compute/using-cloud-tpu-multislice-to-scale-ai-workloads)and best parallel layout are accounted for\(rajbhandari2022deepspeed\)\.

What follows is that the systems\-side recommendations for architectures have largely evaded the scaling laws, taking the model size as given and optimizing the architecture at that size for hardware efficiency\. As an example,anthony2024codesignshow that a transformer at a given scale can be made more efficient on a specific accelerator by aligning its intrinsic dimensions, such as the hidden size, attention\-head count, and vocabulary size, to the structure of the GPU compute kernels, improving throughput at matched accuracy\. A complementary line of work as FlashAttention\(dao2022flashattention\)rearranges the operation arithmetically equivalently and computes the same attention through an IO\-aware, memory\-efficient schedule\. Thus, the training recipe is conventionally built in two disconnected stages, a scaling law stage for the model and data, and a hardware aware tuning stage, as in the Kimi\-K2 report\(team2025kimi\)that treats sparse MoE scaling laws and systems tuning separately\. This separation is consequential at scale, where the model is split across a fleet of GPUs through data, tensor, pipeline, and expert parallelism, as in MegaScale\(jiang2024megascale\)\. The architecture chosen by a scaling law should account for these systems efficiencies rather than model\-FLOPs alone\(bian2025inference\)\.

Our goal in this paper is to fold the systems stage into the scaling law stage, so that a single optimization selects the architecture, the token budget, and the cluster parallelism execution layout together, and we instantiate that framework for sparse MoE language models\. We focus on sparse MoE architectures adopted previously for language model training\(fedus2022switch;sarkar2024revisiting;abnar2025parameters\), where sparsity is obtained by replacing the dense feed\-forward layer of a transformer block with a set ofEEexperts and routing each token to a subset of them\. This decouples the active parametersNactN\_\{\\text\{act\}\}that set per\-token model\-FLOPs from the total stored parametersNtotN\_\{\\text\{tot\}\}, with the expert countEEas one of the knobs controlling total parameter count\. The existing scaling laws for sparse MoE models\(clark2022unified;tian2026towards;krajewski2024scaling\)optimally scale these architecture knobs while ignoring systems efficiency, includingabnar2025parameters, who show that the sparsity factorSS, roughly the proportion ofNtotN\_\{\\text\{tot\}\}that is inactive per token, introduces an additional optimum that classical\(N,D\)\(N,D\)\-only laws cannot see\. We show in Section[6\.1](https://arxiv.org/html/2608.10605#S6.SS1)and Figure[1](https://arxiv.org/html/2608.10605#S1.F1)that a model\-FLOPs budget alone is insufficient to determine a meaningful MoE sparsity\.

A few recent studies begin to surface these systems\-side pressures within the MoE scaling law itself\.ludziejewski2025jointfit a sparse MoE scaling law and find that scaling the number of experts improves loss at fixed model\-FLOPs per token but inflates stored parameters and device\-memory pressure in a way that a model FLOPs budget never charges for, so the practically optimalEEis where the loss surface meets a memory\-aware budget rather than theEEthat minimizes loss at fixedNactN\_\{\\text\{act\}\}\.wan2026holisticpush this further, treating total parameters and model\-FLOPs per token as independent degrees of freedom and arguing that model\-FLOPs per token alone is an inadequate fairness metric for choosing a sparse MoE architecture\.

Building on these studies, we fit a new scaling law whose variables are chosen to expose systems efficiency\. Conventionally, the optimal sparse MoE geometry under a model\-FLOPs budget is derived in two stages, a geometry\-free convex program for the target parameters, sparsity, and data budget, followed by a match to the nearest discrete architecture geometry\. Our scaling analysis shows that sparser MoE models reach lower loss at any model\-FLOPs budget with the optimal sparsity sitting at the boundary, which forces the binding constraint to move from model FLOPs to hardware\-deliverable model FLOPs, the useful compute a cluster can actually deliver over the training window\. We solve a single\-step discrete\-geometry mixed\-integer nonlinear program, which we term MOSAIC\. It minimizes the predicted loss directly over the discrete architecture grid, resolving the boundary sparsity problem and realizing a cluster\-specific sparse MoE geometry as motivated in Figure[1](https://arxiv.org/html/2608.10605#S1.F1)\.

We focus this work on the training side of sparse MoE architecture\-system co\-design, leaving test time scaling laws under hardware costs\(sadhukhan2026kinetics\)as future work\. The closest to our study is the work done insun2026hardwarewhich couple a roofline analysis to architecture selection\. However, roofline alone misses the overheads that flip MoE recommendations, namely expert dispatch and combine, kernel\-launch overhead under fine\-grained experts, all\-to\-all bandwidth as a function of EP topology, and data\-parallel and pipeline\-parallel coupled exposed collective communications\(gale2023megablocks;jiang2024megascale\)\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/x1.png)\(a\)Model FLOPs selects a boundary MoE sparsity\.Under a fixedCmodel=3\.3×1021C\_\{\\mathrm\{model\}\}=3\.3\\times 10^\{21\}, the fitted loss decreases over the calibrated sparsity range, soS⋆=SmaxS^\{\\star\}=S\_\{\\max\}\(left\)\. We do not cap any model dimension, and at each sparsity tillSmaxS\_\{\\max\}\(0\.985 here\), the expert count and activated parameters are optimized under the budget\. Under a hardware budgetCpeakC\_\{\\text\{peak\}\}with 4 AWS p6\-B200 nodes for 5 days, this prescription is infeasible because configurations beyond the feasible boundary ofSS=0\.96 deliver fewerCdeliverableC\_\{\\text\{deliverable\}\}FLOPs \(red region\) than what the fixed prescriptionS⋆S^\{\\star\}requires \(right\)\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/x2.png)\(b\)MOSAIC yields an interior feasible sparsity under the hardware budget\.The hardware\-aware lossℒhw⋆​\(S\)\\mathcal\{L\}^\{\\star\}\_\{\\mathrm\{hw\}\}\(S\)induced byCdeliverableC\_\{\\mathrm\{deliverable\}\}is minimized at the interior sparsityScluster⋆\{\\color\[rgb\]\{0,0\.55,0\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0,0\.55,0\}S^\{\\star\}\_\{\\mathrm\{cluster\}\}\}and is 0\.031 nats lower than atS⋆=Smax\{\\color\[rgb\]\{1,0,0\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{1,0,0\}S^\{\\star\}=S\_\{\\max\}\}\. Each sparsity configuration is evaluated at the token count supported by its own attainable peak MFU under the same 4\-node, 5\-day hardware envelope\.

Figure 1:The systems\-aware scaling motivation for MOSAIC\.##### Contributions\.

To run the workflow of MOSAIC and obtain the optimal sparse MoE scaling dimensions under the constraint of deliverable model FLOPs, we discuss the following contributions and observations:

- •We fit a new four\-dimensional joint MoE pretraining scaling law that expose the systems\-side knobs\. Through our scaling analysis, we show that across model\-FLOPs budgets and within the calibrated sparsity range, the fitted loss optimal sparsity lies at the upper boundary of the data support\. This boundary\-seeking behavior motivates the move from a model\-FLOPs budget to a deliverable model\-FLOPs budget\.
- •We develop an operator\-level analytical performance model that measures the per\-block forward and backward cost of an MoE transformer\. We validate the performance model against training run grids up to 18 billion active parameters, where the mean absolute percentage error of predicted MFU stays under15%15\\%per sweep\.
- •We couple these two components inside MOSAIC and search the architecture grid on a cluster of NVIDIA B200 GPUs under a fixed GPU hour budget\. We show that the lowest\-loss model configuration is not the one that emits the most model FLOPs\. We validate the ranked selections from MOSAIC optimization, on staged pretraining runs of up to250250B total parameters on our infrastructure\. We show that the loss ordering flips from the model\-FLOPs axis to the peak\-equivalent hardware\-compute axis, as MOSAIC predicts\.

We start with describing the abstract formulation of MOSAIC in Section[2](https://arxiv.org/html/2608.10605#S2)\. To apply MOSAIC with sparse MoE model optimization, we first provide a technical overview of the training recipe used in our runs in Section[3](https://arxiv.org/html/2608.10605#S3)\. We then describe the components that instantiate MOSAIC with the sparse MoE architecture: \(1\) sparse MoE scaling laws in Section[4](https://arxiv.org/html/2608.10605#S4)and, \(2\) the performance model in Section[5](https://arxiv.org/html/2608.10605#S5)\. We make the application of MOSAIC for sparse MoEs concrete in Section[6](https://arxiv.org/html/2608.10605#S6), where the scaling law and the performance model come together to show why MoE sparsity has no interior optimum under a model FLOPs budget and why the right budget for a sparse MoE is the hardware\-deliverable FLOPs\. Finally, we discuss the experiments and the results in Section[7](https://arxiv.org/html/2608.10605#S7)\. We provide a snapshot of the most shared symbols used in the rest of the paper in Table[1](https://arxiv.org/html/2608.10605#S1.T1)\.

SymbolMeaningArchitectureLlayersL\_\{\\text\{layers\}\}Number of transformer blocks\.ddResidual\-stream width \(dmodeld\_\{\\text\{model\}\}\)\.dffd\_\{\\text\{ff\}\}Hidden width of the dense FFN that a sparse MoE block replaces\.dexpertd\_\{\\text\{expert\}\}Hidden width of a single expert\.EETotal number of experts per sparse MoE block\.KKTop\-KKactive experts per token\.GGexpert split factor,G=dff/dexpertG=d\_\{\\text\{ff\}\}/d\_\{\\text\{expert\}\}\. Figure[14](https://arxiv.org/html/2608.10605#A2.F14)in Appendix[B](https://arxiv.org/html/2608.10605#A2)visualizes the construction\.NactN\_\{\\text\{act\}\}Active parameters per token \(FLOPs\-relevant\)\.NtotN\_\{\\text\{tot\}\}Total stored parameters of the model \(model\-capacity\)\.SSSparsity,S=1−Nact/NtotS=1\-N\_\{\\text\{act\}\}/N\_\{\\text\{tot\}\}\(fraction of parameters not activated per token\)\.Routing and load balancingH∈ℝT×dH\\in\\mathbb\{R\}^\{T\\times d\}Pre\-MoE hidden states for a batch ofTTtokens \(Algorithm[1](https://arxiv.org/html/2608.10605#alg1)\)\.ϕ∈ℝE\\phi\\in\\mathbb\{R\}^\{E\}Per\-expert routing bias \(LFLB controller\)\.μb\\mu\_\{b\}Bias\-update rate of the LFLB controller\.ρ\\rhoProposal\-to\-selection ratio for Gumbel routing, giving proposal sizeM=min⁡\(ρ​K,E\)M=\\min\(\\rho K,E\)\.si′​\(x\)s^\{\\prime\}\_\{i\}\(x\),si​\(x\)s\_\{i\}\(x\)Clean sigmoid router score and the selection scoresi​\(x\)=si′​\(x\)\+ϕis\_\{i\}\(x\)=s^\{\\prime\}\_\{i\}\(x\)\+\\phi\_\{i\}that adds the balancing bias \(Section[3\.3](https://arxiv.org/html/2608.10605#S3.SS3)\)\.𝒯K​\(x\)\\mathcal\{T\}\_\{K\}\(x\)Set of top\-KKexperts selected for tokenxx\.MOSAIC decision variablesZZArchitecture decision variable of MOSAIC, a realizable sparse MoE geometry with induced scaling law dimensions\(Ntot,S,G\)\(N\_\{\\text\{tot\}\},S,G\)\(Section[2](https://arxiv.org/html/2608.10605#S2)\)\.PℓP\_\{\\ell\}Execution plan,Pℓ=\(PTP,PEP,PCP,PPP,PDP,Bmicro,Ackpt\)P\_\{\\ell\}=\(P\_\{\\text\{TP\}\},P\_\{\\text\{EP\}\},P\_\{\\text\{CP\}\},P\_\{\\text\{PP\}\},P\_\{\\text\{DP\}\},B\_\{\\text\{micro\}\},A\_\{\\text\{ckpt\}\}\), giving the parallelism degrees along tensor \(TP\), expert \(EP\), context \(CP\), pipeline \(PP\), and data parallel \(DP\) groups together with the micro\-batch size and the activation\-checkpointing mode\.BBBatch envelope,B=\(Bglobal,Tseq\)B=\(B\_\{\\text\{global\}\},T\_\{\\text\{seq\}\}\), givingnmb=Bglobal/\(PDP​Bmicro\)n\_\{\\text\{mb\}\}=B\_\{\\text\{global\}\}/\(P\_\{\\text\{DP\}\}\\,B\_\{\\text\{micro\}\}\)microbatches per optimizer step\.DDNumber of training tokens\.CmodelC\_\{\\text\{model\}\}Model FLOPs,Cmodel=6​N​\(Z\)​DC\_\{\\text\{model\}\}=6\\,N\(Z\)\\,D, where the activated parametersN​\(Z\)N\(Z\)areNactN\_\{\\text\{act\}\}for sparse MoEs\.TtrainT\_\{\\text\{train\}\}Wall\-clock training window on the cluster\.CpeakC\_\{\\text\{peak\}\}Raw hardware ceiling,Cpeak=Fpeak​Ndev​TtrainC\_\{\\text\{peak\}\}=F\_\{\\text\{peak\}\}\\,N\_\{\\text\{dev\}\}\\,T\_\{\\text\{train\}\}, whereNdevN\_\{\\text\{dev\}\}is the device count\.FpeakF\_\{\\text\{peak\}\}Per\-device peak throughput\.CdeliverableC\_\{\\text\{deliverable\}\}Deliverable model FLOPs,Cdeliverable​\(Z,Pℓ,B\)=Cpeak⋅MFU​\(Z,Pℓ,B\)⋅ηgood​\(Ndev\)C\_\{\\text\{deliverable\}\}\(Z,P\_\{\\ell\},B\)=C\_\{\\text\{peak\}\}\\cdot\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)\\cdot\\eta\_\{\\text\{good\}\}\(N\_\{\\text\{dev\}\}\), MFU is defined in Equation \([1](https://arxiv.org/html/2608.10605#S2.E1)\) andηgood\\eta\_\{\\text\{good\}\}is the goodput fraction\.Table 1:Symbols used in the paper\.![Refer to caption](https://arxiv.org/html/2608.10605v1/x3.png)Figure 2:The MOSAIC workflow\. It takes the device configuration and the cluster size and training window as inputs, queries the geometry scaling ladder for candidate geometries, the performance model for predicted peak MFU, and the scaling law for predicted loss, and returns the model configuration together with its parallel layout\.

## 2MOSAIC:ModelOptimization viaSystems\-Aware TraIningCo\-design

MOSAIC selects a model architecture, training\-token budget, and distributed execution layout jointly for a fixed cluster and training window\. Its central departure from conventional model FLOPs optimal scaling is that the available model\-FLOPs budget is not treated as an architecture\-independent constant\. In Section[2\.1](https://arxiv.org/html/2608.10605#S2.SS1), we define the resulting architecture\-dependent deliverable budget, we formulate the joint optimization in Section[2\.2](https://arxiv.org/html/2608.10605#S2.SS2), and we instantiate it for sparse MoEs in Section[2\.3](https://arxiv.org/html/2608.10605#S2.SS3)\. Figure[2](https://arxiv.org/html/2608.10605#S1.F2)shows the workflow end to end, from the cluster inputs through the geometry ladder that fixes how the model dimensions co\-scale, the performance model, and the scaling law, to the model configuration and parallel layout it returns\.

### 2\.1From Hardware Budget to Deliverable Model FLOPs

A model FLOPs budget treats every allocation with the sameCmodel​\(Z,D\)=6​N​\(Z\)​DC\_\{\\text\{model\}\}\(Z,D\)=6\\,N\(Z\)\\,Das equally expensive, whereN​\(Z\)N\(Z\)counts the parameters the architecture activates per token \(equals the full parameter count for a dense model\)\. On real clusters, model FLOPs omit a multiplicative systems\-efficiency factor that itself depends on the architecture and its execution plan\. A cluster ofNdevN\_\{\\text\{dev\}\}devices, each with peak per\-device throughputFpeakF\_\{\\text\{peak\}\}, provides the raw hardware ceilingCpeak=Fpeak​Ndev​TtrainC\_\{\\text\{peak\}\}=F\_\{\\text\{peak\}\}\\,N\_\{\\text\{dev\}\}\\,T\_\{\\text\{train\}\}over a training windowTtrainT\_\{\\text\{train\}\}\. It is the physical maximum available over that window and only a fraction of this ceiling becomes useful model computation\.

Before discussing how we define useful model computation, we define Model FLOPs Utilization \(MFU\) as is used in our work\. It is the fraction of the cluster’s peak throughput that a training run converts into the FLOPs mathematically required by the model\(chowdhery2023palm\)\. An execution planPℓP\_\{\\ell\}specifies how the model is sharded and its computation scheduled across the GPU fleet\(meta2025llama3isca\)\. For architectureZZ, execution planPℓP\_\{\\ell\}, and batch envelopeBB, lettiter​\(Z,Pℓ,B\)t\_\{\\text\{iter\}\}\(Z,P\_\{\\ell\},B\)denote the wall\-clock time of one optimizer step processingBglobal​TseqB\_\{\\text\{global\}\}\\,T\_\{\\text\{seq\}\}tokens, whereTseqT\_\{\\text\{seq\}\}is the sequence length\. We define MFU as

MFU​\(Z,Pℓ,B\)=6​N​\(Z\)​Bglobal​Tseqtiter​\(Z,Pℓ,B\)​Ndev​Fpeak\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)\\;=\\;\\frac\{6\\,N\(Z\)\\;B\_\{\\text\{global\}\}\\,T\_\{\\text\{seq\}\}\}\{t\_\{\\text\{iter\}\}\(Z,P\_\{\\ell\},B\)\\;N\_\{\\text\{dev\}\}\\;F\_\{\\text\{peak\}\}\}\(1\)It is the model FLOPs per token executed per second, normalized by the peak FLOPs rateFpeak​NdevF\_\{\\text\{peak\}\}N\_\{\\text\{dev\}\}the cluster could sustain\. The numerator counts only the FLOPs the model mathematically requires, so activation recomputation, though executed on the hardware, earns no credit\. This distinguishes MFU from hardware FLOPs utilization \(HFU\), and MFU is the metric that tracks time\-to\-loss\. MFU captures per\-step execution efficiency, including kernel quality, compute and communication overlap, pipeline bubbles, and expert\-parallel all\-to\-all\. The goodput factorηgood​\(Ndev\)\\eta\_\{\\text\{good\}\}\(N\_\{\\text\{dev\}\}\)captures the fraction of allocated wall\-clock spent training rather than checkpointing, recovering from failures, or waiting on the input pipeline\. We define the deliverable model FLOPs as

Cdeliverable​\(Z,Pℓ,B\)=Cpeak⋅MFU​\(Z,Pℓ,B\)⋅ηgood​\(Ndev\)\.C\_\{\\text\{deliverable\}\}\(Z,P\_\{\\ell\},B\)\\;=\\;C\_\{\\text\{peak\}\}\\;\\cdot\\;\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)\\;\\cdot\\;\\eta\_\{\\text\{good\}\}\(N\_\{\\text\{dev\}\}\)\.\(2\)Unlike a conventional compute budget,CdeliverableC\_\{\\text\{deliverable\}\}depends on the architecture and the layoutPℓP\_\{\\ell\}used to execute it\. A candidate training configuration\(Z,D\)\(Z,D\)is feasible to train on the cluster within the windowTtrainT\_\{\\text\{train\}\}only if it satisfies the constraint

Cmodel​\(Z,D\)≤maxPℓ∈𝒫feas​\(Z,B\)⁡Cdeliverable​\(Z,Pℓ,B\)C\_\{\\text\{model\}\}\(Z,D\)\\;\\leq\\;\\max\_\{P\_\{\\ell\}\\in\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)\}\\,C\_\{\\text\{deliverable\}\}\(Z,P\_\{\\ell\},B\)\(3\)where𝒫feas​\(Z,B\)\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)collects the execution plans that fit the cluster constraint, defined in Section[2\.2](https://arxiv.org/html/2608.10605#S2.SS2)\.

### 2\.2Joint Architecture\-Systems Optimization

MOSAIC formalizes hardware\-aware model optimization as a mixed\-integer nonlinear program \(MINLP\) that couples the architecture to the systems layout it will execute under\. It optimizes three coupled choices, the architectureZZ, drawn from a finite set𝒵disc\\mathcal\{Z\}\_\{\\text\{disc\}\}of realizable architectures, each carrying the scaling law dimensions its geometry induces, the training tokensDD, and the execution planPℓP\_\{\\ell\}of Section[2\.1](https://arxiv.org/html/2608.10605#S2.SS1)\. The hardware inputs are\(Fpeak,Ndev,Ttrain,Mdev\)\(F\_\{\\text\{peak\}\},N\_\{\\text\{dev\}\},T\_\{\\text\{train\}\},M\_\{\\text\{dev\}\}\), whereMdevM\_\{\\text\{dev\}\}is the per\-device memory cap, and the batch envelopeB=\(Bglobal,Tseq\)B=\(B\_\{\\text\{global\}\},T\_\{\\text\{seq\}\}\)fixes the tokens per optimizer step\. The execution plans feasible for an architecture are collected in

𝒫feas​\(Z,B\)=\{Pℓ∈𝒫disc​\(Ndev\):Mem​\(Z,Pℓ,B\)≤Mdev\}\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)\\;=\\;\\bigl\\\{\\,P\_\{\\ell\}\\in\\mathcal\{P\}\_\{\\text\{disc\}\}\(N\_\{\\text\{dev\}\}\)\\;:\\;\\mathrm\{Mem\}\(Z,P\_\{\\ell\},B\)\\leq M\_\{\\text\{dev\}\}\\,\\bigr\\\}\(4\)where𝒫disc​\(Ndev\)\\mathcal\{P\}\_\{\\text\{disc\}\}\(N\_\{\\text\{dev\}\}\)is the finite set of realizable execution plans the target training stack exposes onNdevN\_\{\\text\{dev\}\}devices, comprising the parallelism degrees that tile the cluster, the micro\-batch sizeBmicroB\_\{\\text\{micro\}\}, and the activation\-checkpointing modeAckptA\_\{\\text\{ckpt\}\}\. Over this feasible plan set, we solve the following optimization

minZ∈𝒵disc,D\\displaystyle\\min\_\{Z\\in\\mathcal\{Z\}\_\{\\text\{disc\}\},\\;D\}\\quadℒ​\(Z,D\)\\displaystyle\\mathcal\{L\}\\bigl\(Z,D\\bigr\)\(5\)subject toCmodel​\(Z,D\)≤maxPℓ∈𝒫feas​\(Z,B\)⁡Cdeliverable​\(Z,Pℓ,B\)\.\\displaystyle C\_\{\\text\{model\}\}\(Z,D\)\\;\\leq\\;\\max\_\{P\_\{\\ell\}\\in\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)\}\\,C\_\{\\text\{deliverable\}\}\(Z,P\_\{\\ell\},B\)\.\(C\-FLOPs\)

The objectiveℒ\\mathcal\{L\}is the scaling law, which predicts loss from the architecture specific dimensionZZand the dataDD\. The loss functional form evaluatesZZthrough the scaling law dimensions it induces, whereasMFU\(\.\)\\mathrm\{MFU\}\(\.\)andMem\(\.\)\\mathrm\{Mem\}\(\.\)evaluate the concrete architecture itself, so all three are well\-defined functions ofZZeven though several architectures can share the same dimensions\. The execution plan enters only through the inner maximum, so the outer problem optimizes\(Z,D\)\(Z,D\)alone\. All systems constraints sit inside𝒫feas​\(Z,B\)\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\), which means that the maximum never returns an infeasible plan\. Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\) is the feasibility condition of Equation \([3](https://arxiv.org/html/2608.10605#S2.E3)\)\. Equivalently, MOSAIC chooses for each architecture the best feasible plan as

Pℓ⋆​\(Z\)=arg⁡maxPℓ∈𝒫feas​\(Z,B\)⁡Cdeliverable​\(Z,Pℓ,B\)\.P\_\{\\ell\}^\{\\star\}\(Z\)\\;=\\;\\arg\\max\_\{P\_\{\\ell\}\\in\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)\}\\;C\_\{\\text\{deliverable\}\}\(Z,P\_\{\\ell\},B\)\.\(6\)
We note that Equation \([5](https://arxiv.org/html/2608.10605#S2.E5)\) is a bilevel non\-convex mixed\-integer nonlinear program, an optimization problem that is bilevel because of the embedded maximization over𝒫feas\\mathcal\{P\}\_\{\\text\{feas\}\}in Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\)\. It is mixed\-integer nonlinear since the layout degrees are integer and the architecture is drawn from the discrete set𝒵disc\\mathcal\{Z\}\_\{\\text\{disc\}\}\. The constraint functionsMFU​\(⋅\)\\mathrm\{MFU\}\(\\cdot\),Mem​\(⋅\)\\mathrm\{Mem\}\(\\cdot\)are nonlinear, and non\-convex because of the discrete variables of the realizable plan set𝒫disc​\(Ndev\)\\mathcal\{P\}\_\{\\text\{disc\}\}\(N\_\{\\text\{dev\}\}\)together with the non\-convexity of MFU and memory inPℓP\_\{\\ell\}\. We collect in Appendix[D\.2](https://arxiv.org/html/2608.10605#A4.SS2)three structural properties of this program, including a closed\-form elimination ofDD, that make its exact solution by enumeration tractable\.

### 2\.3Instantiating MOSAIC for Sparse MoEs

We put Equation \([5](https://arxiv.org/html/2608.10605#S2.E5)\) to be architecture\-agnostic\. Applying it requires three components, a realizable architecture space induced byZZ, a predictive loss functionℒ​\(Z,D\)\\mathcal\{L\}\(Z,D\), and predictors forMFU​\(Z,Pℓ,B\)\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)andMem​\(Z,Pℓ,B\)\\mathrm\{Mem\}\(Z,P\_\{\\ell\},B\)\. We define the sparse MoE realization ofZZin Section[3](https://arxiv.org/html/2608.10605#S3), under which𝒵disc\\mathcal\{Z\}\_\{\\mathrm\{disc\}\}becomes the finite set of realizable sparse MoE geometries and the activated parameter count becomesN​\(Z\)=NactN\(Z\)=N\_\{\\mathrm\{act\}\}\. The execution plan additionally includes expert parallelism, and in our Megatron\-Core realization, expert parallelism repartitions the data\-parallel domain rather than introducing an independent world\-size factor\. We give the exact mesh factorization and divisibility constraints used to construct𝒫disc​\(Ndev\)\\mathcal\{P\}\_\{\\text\{disc\}\}\(N\_\{\\text\{dev\}\}\)for our MoE models in Appendix[D\.1](https://arxiv.org/html/2608.10605#A4.SS1)\. We discuss each of these three components applied to sparse MoEs in the following sections\. Figure[1](https://arxiv.org/html/2608.10605#S1.F1)instantiates the constraint of Equation[3](https://arxiv.org/html/2608.10605#S2.E3)with real numbers where aCmodel=3\.3×1021C\_\{\\mathrm\{model\}\}=3\.3\\times 10^\{21\}prescription on 4 p6\-B200 nodes for 5 days becomes infeasible beyondSfeasmax=0\.96S\_\{\\mathrm\{feas\}\}^\{\\max\}=0\.96\. Re\-optimizing under the same hardware envelope yields an interior cluster\-optimal sparsity whose predicted loss is 0\.031 nats lower than the model FLOPs optimal boundary configuration\. These constraints let us treat the model dimensions as decision variables of the optimization rather than quantities capped in advance, which we argue is the more natural formulation for MoE\-like models\.

## 3Sparse MoE Pretraining: Scaling Recipe

In this section, we discuss the dimensionsZZfor sparse MoEs in the context of scaling laws and describe the training recipe used to collect the scaling law data\. We lay out the details of the architecture that is shared by all the runs in our study\. Our architecture and training recipe follows the broader work from existing sparse MoE scaling studies ofkrajewski2024scalingandabnar2025parameterswith some adjustments, and we describe those differences in the following sections\.

### 3\.1Training Setup

We experiment with decoder\-only transformer language models in which every feed\-forward layer is replaced by an MoE layer\. The first transformer block is kept dense and all remaining blocks are sparse MoE layers\. Following the architecture design in DeepSeek\-V3\(deepseekai2024v3;dai2024deepseekmoe\), every MoE layer in the transformer block carries one shared expert that is always active in addition to theKKrouted experts\. Each training run is optimized for text\-to\-text with the next\-token\-prediction language modeling objective and all runs use bf16 weights and activations\. We focus on the MoE layer in this section since that forms the basis of the scaling laws in the rest of the paper\. We use the standard parametrization recipe for weight initializations and we adopt AdamW optimizer\(loshchilov2017decoupled\)for all the runs\. Similar to the Chinchilla training runs\(hoffmann2022training\), we consider a single pass over the tokens and we use the smoothed training loss of the final checkpoint of each run for the scaling law fits\. The rest of the technical details of the data, training recipe and the hyper\-parameter choices including the optimizer choices and learning rate scheduler are described in Appendix[A\.1](https://arxiv.org/html/2608.10605#A1.SS1)\.

### 3\.2MoE Layer

We useLlayersL\_\{\\text\{layers\}\}to denote the number of transformer blocks,d≡dembdd\\\!\\equiv\\\!d\_\{\\text\{embd\}\}as the residual stream width, anddffd\_\{\\text\{ff\}\}as the hidden width of the corresponding dense feed\-forward layer that a sparse MoE layer replicates intoEEcopies333dffd\_\{\\text\{ff\}\}is conventionally a fixed multiple of the residual widthdd\. For example, GPT\-2\(radford2019language\)usesdff=4​dd\_\{\\text\{ff\}\}=4d\.\. Each MoE layer consists ofEEexperts\{fi\}i=1E\\\{f\_\{i\}\\\}\_\{i=1\}^\{E\}, where every expert is itself a two\-layer feed\-forward network of the formfi​\(x\)=Wi\(2\)​σ​\(Wi\(1\)​x\)f\_\{i\}\(x\)=W^\{\(2\)\}\_\{i\}\\,\\sigma\\\!\\left\(W^\{\(1\)\}\_\{i\}x\\right\)withWi\(1\)∈ℝdexpert×dW^\{\(1\)\}\_\{i\}\\in\\mathbb\{R\}^\{d\_\{\\text\{expert\}\}\\times d\},Wi\(2\)∈ℝd×dexpertW^\{\(2\)\}\_\{i\}\\in\\mathbb\{R\}^\{d\\times d\_\{\\text\{expert\}\}\}, andσ​\(⋅\)\\sigma\(\\cdot\)a gated activation, in our runs being SwiGLU\. We keepdffd\_\{\\text\{ff\}\}anddexpertd\_\{\\text\{expert\}\}as separate symbols because a sparse MoE layer can split the reference FFN of widthdffd\_\{\\text\{ff\}\}into narrower experts of widthdexpertd\_\{\\text\{expert\}\}, so the two coincide only when each expert is a full\-width dense FFN, something we will explore in Section[4](https://arxiv.org/html/2608.10605#S4)\. We follow the Switch\-style token\-choice routing formulation\(fedus2022switch;liu2023sparsemixer\)\. A learnable routerWr∈ℝE×dW\_\{r\}\\in\\mathbb\{R\}^\{E\\times d\}produces a per\-token score vector𝒔​\(x\)∈ℝE\\bm\{s\}\(x\)\\in\\mathbb\{R\}^\{E\}, and the top\-KKexperts under𝒔​\(x\)\\bm\{s\}\(x\)are selected for each token in a forward pass\. The sparse MoE layer output is the gating\-weighted sum of their outputs,

y​\(x\)=∑i∈𝒯K​\(x\)π~i​\(x\)​fi​\(x\),𝒯K​\(x\)=Top⁡K​\(𝒔​\(x\)\)y\(x\)\\;=\\;\\sum\_\{i\\,\\in\\,\\mathcal\{T\}\_\{K\}\(x\)\}\\tilde\{\\pi\}\_\{i\}\(x\)\\,f\_\{i\}\(x\),\\qquad\\mathcal\{T\}\_\{K\}\(x\)\\;=\\;\\operatorname\{Top\}\\\!K\\bigl\(\\bm\{s\}\(x\)\\bigr\)\(7\)whereπ~i​\(x\)\\tilde\{\\pi\}\_\{i\}\(x\)are the probabilities the selected set𝒯K​\(x\)\\mathcal\{T\}\_\{K\}\(x\)\. In the following section, we describe the routing operations that result in theTopK\(\.\)\\operatorname\{Top\}\\\!K\(\.\)outputs and the associated probabilitiesπ~i​\(x\)\\tilde\{\\pi\}\_\{i\}\(x\)and we treat Equation[7](https://arxiv.org/html/2608.10605#S3.E7)as the canonical sparse MoE layer definition\.

Algorithm 1Sparse MoE token\-choice routing for a single MoE layer, batch ofTTtokens routed toEEexperts0:Pre\-MoE hidden states

H∈ℝT×dH\\in\\mathbb\{R\}^\{T\\times d\}\(attention output of the current block added to the previous\-layer residual stream\); router weights

Wr∈ℝE×dW\_\{r\}\\in\\mathbb\{R\}^\{E\\times d\}; per\-expert bias

ϕ∈ℝE\\phi\\in\\mathbb\{R\}^\{E\}; hyperparameters

KK\(top\-

KK\), Gumbel proposal size

M=min⁡\(ρ​K,E\)M=\\min\(\\rho K,E\), Gumbel temperature

τ\\tau\.

0:Top\-

KKexpert indices

𝒯K∈\{1,…,E\}T×K\\mathcal\{T\}\_\{K\}\\in\\\{1,\\dots,E\\\}^\{T\\times K\}, where row

𝒯K​\(xt\)=Top⁡K​\(𝒔​\(xt\)\)\\mathcal\{T\}\_\{K\}\(x\_\{t\}\)=\\operatorname\{Top\}\\\!K\(\\bm\{s\}\(x\_\{t\}\)\)over the proposal set

𝒞t\\mathcal\{C\}\_\{t\}is the selected set for token

ttin Eq\.[7](https://arxiv.org/html/2608.10605#S3.E7); normalized gate probabilities

π~∈ℝ≥0T×K\\tilde\{\\pi\}\\in\\mathbb\{R\}\_\{\\geq 0\}^\{T\\times K\}with

∑i∈𝒯K​\(xt\)π~t,i=1\\sum\_\{i\\in\\mathcal\{T\}\_\{K\}\(x\_\{t\}\)\}\\tilde\{\\pi\}\_\{t,i\}=1; per\-expert token counts

n∈ℕEn\\in\\mathbb\{N\}^\{E\}\. Throughout,

i∈\{1,…,E\}i\\in\\\{1,\\dots,E\\\}identifies an expert and

t∈\{1,…,T\}t\\in\\\{1,\\dots,T\\\}a token\.

1:

W~r←RowNorm​\(Wr,cr\)\\tilde\{W\}\_\{r\}\\leftarrow\\textsc\{RowNorm\}\(W\_\{r\},c\_\{r\}\)\{

ℓ2\\ell\_\{2\}\-normalize each router row to target norm

crc\_\{r\}\}

2:

Θ←H​W~r⊤\\Theta\\leftarrow H\\,\\tilde\{W\}\_\{r\}^\{\\\!\\top\}\{router logits\}

3:

s′←σ​\(Θ\)s^\{\\prime\}\\leftarrow\\sigma\(\\Theta\)\{clean scores, no bias\}

4:

s←s′\+ϕs\\leftarrow s^\{\\prime\}\+\\phi\{LFLB bias\-added selection scores, Eq\.[18](https://arxiv.org/html/2608.10605#A1.E18)in Appendix[A\.3](https://arxiv.org/html/2608.10605#A1.SS3)\}

5:foreach token

t∈Tt\\in T\(parallel\)do

6:

gt,i∼Gumbel​\(0,1\)g\_\{t,i\}\\sim\\mathrm\{Gumbel\}\(0,1\)for

i=1,…,Ei=1,\\dots,E\{i\.i\.d\. standard Gumbel per expert\}

7:

gt←gt−1T​∑t′=1Tgt′g\_\{t\}\\leftarrow g\_\{t\}\-\\frac\{1\}\{T\}\\\!\\sum\_\{t^\{\\prime\}=1\}^\{T\}g\_\{t^\{\\prime\}\}\{per\-expert noise centering\}

8:

𝒞t←Top−⁡Mi​\(log⁡st,i′\+τ​gt,i\)\\mathcal\{C\}\_\{t\}\\leftarrow\\operatorname\*\{Top\-\}\\\!M\_\{i\}\\bigl\(\\log s^\{\\prime\}\_\{t,i\}\+\\tau\\,g\_\{t,i\}\\bigr\)\{proposal on clean scores\}

9:

𝒯K,t←Top−⁡Ki∈𝒞t​\(st,i\)\\mathcal\{T\}\_\{K,t\}\\leftarrow\\operatorname\*\{Top\-\}\\\!K\_\{i\\in\\mathcal\{C\}\_\{t\}\}\\bigl\(s\_\{t,i\}\\bigr\)\{bias\-aware correction\}

10:

ζt←∑j∈𝒯K,tst,j′\\zeta\_\{t\}\\leftarrow\\sum\_\{j\\in\\mathcal\{T\}\_\{K,t\}\}s^\{\\prime\}\_\{t,j\}
11:

π~t,i←st,i′ζt,i∈𝒯K,t\\tilde\{\\pi\}\_\{t,i\}\\leftarrow\\frac\{s^\{\\prime\}\_\{t,i\}\}\{\\zeta\_\{t\}\},\\quad i\\in\\mathcal\{T\}\_\{K,t\}\{gates use clean scores,

∑=1\\sum=1\}

12:endfor

13:

ni←∑t=1T𝟙​\[i∈𝒯K,t\],i=1,…,En\_\{i\}\\leftarrow\\sum\_\{t=1\}^\{T\}\\mathbb\{1\}\\\!\\left\[i\\in\\mathcal\{T\}\_\{K,t\}\\right\],\\quad i=1,\\dots,E\{entries of the per\-expert token\-count vector

nn\}

14:return

𝒯K,π~,n\\mathcal\{T\}\_\{K\},\\ \\tilde\{\\pi\},\\ n

### 3\.3Routing and Load Balancing

Stable large\-scale sparse MoE training hinges on the routing recipe that selects the active experts per token, since imbalances in token\-to\-expert assignment compound into load\-balancing issues\(hu2025sigma;wang2024lossfree\), loss spikes, and a degraded effective parameter count\. We describe the routing choices that kept training stable as we scaled up, a prerequisite for reliable scaling law extrapolation\. There have been myriad studies focused on designing such routing algorithms that improve both the training dynamics and the end model accuracy\(muqeeth2023soft;zoph2022st;zhou2022mixture;hu2025sigma\)\. We build on the auxiliary\-loss\-free load balancing \(LFLB\) strategy ofwang2024lossfree, in which a per\-expert bias is updated online to counteract observed load imbalance, but augment it with several modifications that we found necessary for stable training across our scaling sweeps\.

In its original form, LFLB scores each token\-expert pair by a sigmoid affinitysi′​\(x\)=σ​\(θi​\(x\)\)s^\{\\prime\}\_\{i\}\(x\)=\\sigma\(\\theta\_\{i\}\(x\)\)withiidenoting an expert andθi​\(x\)=w~r,i⊤​x\\theta\_\{i\}\(x\)=\\tilde\{w\}\_\{r,i\}^\{\\\!\\top\}xdenoting the router logit\. It selects the top\-KKexperts by a bias\-adjusted scoresi​\(x\)=si′​\(x\)\+ϕis\_\{i\}\(x\)=s^\{\\prime\}\_\{i\}\(x\)\+\\phi\_\{i\}, and combines their outputs with the clean scoressi′​\(x\)s^\{\\prime\}\_\{i\}\(x\)renormalized over the selected set\(wang2024lossfree;deepseekai2024v3\)\. The per\-expert biasϕi\\phi\_\{i\}is a controller state nudged once per training step toward under\-used experts from the observed per\-expert token countsnin\_\{i\}, so the bias fixed at the end of one step steers selection on the next\. We keep this sigmoid\-plus\-bias score \(Eq\.[18](https://arxiv.org/html/2608.10605#A1.E18)in Appendix[A\.3](https://arxiv.org/html/2608.10605#A1.SS3)\) but adjust the selection rule and the bias controller, as summarized next\.

We introduce a few additions to the base LFLB recipe, laid out in Algorithm[1](https://arxiv.org/html/2608.10605#alg1)and detailed in Appendix[A\.3](https://arxiv.org/html/2608.10605#A1.SS3)\. The algorithm takes as input a batch ofTTtokens with their pre\-MoE hidden statesHH, the router weightsWrW\_\{r\}, and the current per\-expert biasϕ\\phi, and returns the top\-KKexpert indices𝒯K\\mathcal\{T\}\_\{K\}selected per token, their normalized gate probabilitiesπ~\\tilde\{\\pi\}, and the per\-expert token countsnnthat drive the bias update for the next step\. The main additions to LFLB are as follows\.

- •The expert bias is updated online by an RMS\-normalized imbalance controller and re\-centered each step, so the correction tracks the shape of the imbalance rather than the batch size\.
- •The top\-KKselection involves a Gumbel proposal\-correction sampler\. It entails a stochastic top\-MMproposal on the clean sigmoid scores followed by a bias\-aware top\-KKcorrection \(M\>KM\>K, withMMa pre\-selected hyperparameter fixed across all runs\), with the per\-expert noise centered across the batch\. The top\-MMproposal followed by the top\-KKcorrection separates exploration from final routing\. We use Gumbel scores only to expand the candidate set toMM, then choose the finalKKusing clean scores inside that set\. Since applying Gumbel noise over all experts can cause large arbitrary flips, we restrict it to a top\-MMcandidate pool amongEEexperts\. It lets near\-boundary experts get exposure while preventing far\-tail experts from winning due to noise\. Lines 5–9 of Algorithm[1](https://arxiv.org/html/2608.10605#alg1)follow this procedure\.
- •The router rows areℓ2\\ell\_\{2\}\-normalized every forward pass so routing depends on direction rather than norm drift\. Load balancing is thus carried by the bias controller and the noise, with no batch\-level load\-balancing loss added to the objective\. We do retain the sequence\-level auxiliary loss ofdeepseekai2024v3, described in Appendix[A\.3](https://arxiv.org/html/2608.10605#A1.SS3)\.

Together, these changes provide a stable routing recipe across the sweep without a batch\-level load\-balancing loss\. That stability is what lets us treat the final\-checkpoint loss of every run as a comparable fit target for the scaling laws of Section[4](https://arxiv.org/html/2608.10605#S4)\.

### 3\.4SparsitySSand Expert Split FactorGG

In considering the geometry of MoE models, the two architectural dimensions of interest are the total expert countEEand the per\-token activated expertsKK\. Together withLlayersL\_\{\\text\{layers\}\},dd, anddffd\_\{\\text\{ff\}\}, these determine an active parameter budgetNactN\_\{\\text\{act\}\}\(the model FLOPs\-relevant parameter count seen by any individual token\) and a total parameter budgetNtotN\_\{\\text\{tot\}\}\(the memory footprint of the model\)\. In this section, we discuss two architectural knobs, the sparsitySSand the expert split factorGG, that we later use to parametrize the sparse MoE scaling laws\. We recall from Table[1](https://arxiv.org/html/2608.10605#S1.T1)that the two parameter budgets set the sparsityS=1−Nact/NtotS=1\-N\_\{\\text\{act\}\}/N\_\{\\text\{tot\}\}, the fraction of parameters left inactive per token\. To compare architectures whose experts differ in size, we useexpert split factordefined as:

G≔dffdexpertG\\;\\;\\coloneqq\\;\\;\\frac\{d\_\{\\text\{ff\}\}\}\{d\_\{\\text\{expert\}\}\}\(8\)
It is the ratio of the dense FFN widthdffd\_\{\\text\{ff\}\}that an MoE block replaces to a single expert’s hidden widthdexpertd\_\{\\text\{expert\}\}, measuring how finely that dense FFN has been partitioned into experts\. We visually show the construction of this expert split factorGGin an MoE layer in Figure[14](https://arxiv.org/html/2608.10605#A2.F14)in Appendix[B](https://arxiv.org/html/2608.10605#A2)\.dffd\_\{\\text\{ff\}\}is the intermediate \(hidden\) width of the reference full FFN, a fixed multiple of the residual\-stream widthdmodeld\_\{\\text\{model\}\}, and is distinct fromdmodeld\_\{\\text\{model\}\}itself anddexpert=dffGd\_\{\\text\{expert\}\}=\\frac\{d\_\{\\text\{ff\}\}\}\{G\}is the per\-expert slice of it\.

TheG=1G\\\!=\\\!1case recovers the setting studied inclark2022unifiedin which each expert is itself a full\-width FFN\.G\>1G\\\!\>\\\!1corresponds to fine\-grained experts whose width shrinks proportionally as their count grows\. Equivalently, at a fixed total parameter count \(or fixed total expert width\), a higherGGmeans smaller, but more finer\-grained experts in the MoE layer\. The ratioEG\\frac\{E\}\{G\}is the effective expansion rate of total to dense FFN parameters used bykrajewski2024scaling, so increasingGGat fixedEG\\frac\{E\}\{G\}yields more, smaller experts at constant total parameter count\.

## 4Sparse MoE Scaling Lawsℒ\(\.\)\\mathcal\{L\}\(\.\)

Building on the sparse MoE architecture discussed in Section[3](https://arxiv.org/html/2608.10605#S3), we elaborate on whatZZentails and we lay out the desiderata behind the scaling law formulations we target in our study\. We want a loss law whereZZmeets three requirements\. First, a concrete discrete MoE geometry must map to the law’s dimensions, so a downstream optimizer can realize the fitted optimum as a real architecture rather than an abstract parameter count, which is why we parameterize in total parameters and geometry knobs\. Second, those dimensional knobs must expose how the model geometry acts on the systems side including communication, and memory and not on loss alone\. Third, we prefer that the law stay low\-dimensional, because if each of theuufree dimensions in a scaling law is evaluated atvvdiscrete values using a full factorial design, the number of configurations scales asO​\(vu\)O\(v^\{u\}\)\. The laws fitted in the study byabnar2025parametersand bytian2026towardsare on essentially the same axes, differing mainly in the shared\-expert ratio\. We build on both while adding the geometry dimension the systems side needs\. We utilize the sparsitySSand expert split factorGGdescribed in Section[3](https://arxiv.org/html/2608.10605#S3), and argue thatGGis a necessary axis beyond those considered inabnar2025parameters\. We then review theabnar2025parameterslaw on our training runs in Section[4\.1](https://arxiv.org/html/2608.10605#S4.SS1), following which we discuss the final scaling function form we use for the optimization in Section[4\.2](https://arxiv.org/html/2608.10605#S4.SS2)\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_G4_main.png)\(a\)G=4G=4\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_G8.png)\(b\)G=8G=8\.

Figure 3:Predicted\-vs\-observed training loss for theℒ​\(Ntot,S,D\)\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D\)fit in Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\), fit per expert split factor\. The panels for every stratum we fit,G∈\{1,2,4,8\}G\\\!\\in\\\!\\\{1,2,4,8\\\}, are in Figure[15](https://arxiv.org/html/2608.10605#A2.F15)of Appendix[B](https://arxiv.org/html/2608.10605#A2)\.### 4\.1Scaling Laws with Sparsity

Before fitting the functional form following the desiderata above, we follow the functional form studied inabnar2025parameters, and parameterize per\-token training loss as a separable Chinchilla\-style law in total parameters and tokens, with an additional sparsity term, so the induced scaling law dimensions are\(Ntot,S\)\(N\_\{\\text\{tot\}\},S\),

ℒ​\(Ntot,S,D\)=aNtotα⏟capacity\+bDβ⏟data\+c\(1−S\)λ⏟sparsity floor\+j\(1−S\)δ​Ntotγ⏟interaction\+e\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D\)\\;=\\;\\underbrace\{\\frac\{a\}\{N\_\{\\text\{tot\}\}^\{\\alpha\}\}\}\_\{\\text\{capacity\}\}\\;\+\\;\\underbrace\{\\frac\{b\}\{D^\{\\beta\}\}\}\_\{\\text\{data\}\}\\;\+\\;\\underbrace\{\\frac\{c\}\{\(1\-S\)^\{\\lambda\}\}\}\_\{\\text\{sparsity floor\}\}\\;\+\\;\\underbrace\{\\frac\{j\}\{\(1\-S\)^\{\\delta\}\\,N\_\{\\text\{tot\}\}^\{\\gamma\}\}\}\_\{\\text\{interaction\}\}\\;\+\\;e\(9\)with sparsitySS, which differs from how sparsity \(=1−KE=1\-\\frac\{K\}\{E\}\) is defined in that study\. The first two terms recover the Chinchilla form\(hoffmann2022training\)in the limitS→0S\\\!\\to\\\!0\(dense model\)\.

We fit the functional form in Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\) separately for each expert split factorGG, holding out the largest models by active parameters to test upward extrapolation\. The fits are statistically tight\. The predicted and observed loss collapse onto the identity line with held\-out RMSE between0\.00020\.0002and0\.00370\.0037in loss units across the four expert split factors, under half a percent relative error\. Figure[3](https://arxiv.org/html/2608.10605#S4.F3)showsG=4G\\\!=\\\!4andG=8G\\\!=\\\!8, and the fullG∈\{1,2,4,8\}G\\\!\\in\\\!\\\{1,2,4,8\\\}panel is given in Appendix[B\.1](https://arxiv.org/html/2608.10605#A2.SS1)\. What the parity plots also reveal is that the fitted coefficients change substantially acrossGG, so no single set of\(Ntot,S,D\)\(N\_\{\\text\{tot\}\},S,D\)coefficients describes all the expert split factors at once\. This is what leads us to buildGGinto the law in Section[4\.2](https://arxiv.org/html/2608.10605#S4.SS2)\. In Figure[4](https://arxiv.org/html/2608.10605#S4.F4), we find that three trends are monotone overS∈\[0\.5,0\.981\]S\\in\[0\.5,0\.981\]\. These three trends agree with thoseabnar2025parametersreport, recovered here on our own runs and under our sparsity definition\. We return to this figure and its consequence for the model\-FLOPs optimum in Section[6\.1](https://arxiv.org/html/2608.10605#S6.SS1)\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/x4.png)Figure 4:Model FLOPs optimalNtot⋆N\_\{\\text\{tot\}\}^\{\\star\},Nact⋆N\_\{\\text\{act\}\}^\{\\star\}, andℒ⋆\\mathcal\{L\}^\{\\star\}as a function of sparsitySS, at five budgetsCC\. The curves are computed from our ownG=4G\\\!=\\\!4fit of Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\) on our runs\. Larger MoE sparsitySSincreases optimal total parameters, decreases optimal active parameters, and lowers achievable loss at every model\-FLOPs budget\.
### 4\.2GG\-augmented Joint Scaling Law

Before adding the expert split factorGGto the law, we note thatGGand the sparsitySSare distinct knobs that happen to share the parameter counts\. The two move different geometry variables and appear coupled only because both these dimensions modulateNtotN\_\{\\text\{tot\}\}andNactN\_\{\\text\{act\}\}\. Empirically they are close to orthogonal on our grid, where the Pearson correlation betweenlog⁡G\\log Gand each oflog⁡Ntot\\log N\_\{\\text\{tot\}\},SS, andlog⁡Nact\\log N\_\{\\text\{act\}\}is at most0\.020\.02in magnitude, because the sweep varieddexpertd\_\{\\text\{expert\}\}independently of\(E,K\)\(E,K\)and everyG∈\{1,2,4,8\}G\\\!\\in\\\!\\\{1,2,4,8\\\}is observed across nearly the fullNtotN\_\{\\text\{tot\}\}andSSranges\.

A starting point to understand how to augment the scaling functional form with the expert split factorGGis the conventional model FLOPs optimal recipe\. Figure[5](https://arxiv.org/html/2608.10605#S4.F5)shows that the loss curves inNtotN\_\{\\text\{tot\}\}fan out rather than stay parallel asGGorSSvaries, so the effect of the expert split factor and of sparsity is not a rigid shift but depends onNtotN\_\{\\text\{tot\}\}, which is a coupling a separable law in\(Ntot,S,D\)\(N\_\{\\text\{tot\}\},S,D\)cannot express\. Following this observation, we tieGG,SS, andNtotN\_\{\\text\{tot\}\}together in a single interaction term\. The joint scaling law we use, on the scaling law dimensions\(Ntot,S,G\)\(N\_\{\\text\{tot\}\},S,G\)induced byZZ, is

ℒ​\(Ntot,S,D,G\)\\displaystyle\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D,G\)=aNtotα\+bDβ\+c\(1−S\)λ\\displaystyle\\;=\\;\\frac\{a\}\{N\_\{\\text\{tot\}\}^\{\\alpha\}\}\\;\+\\;\\frac\{b\}\{D^\{\\beta\}\}\\;\+\\;\\frac\{c\}\{\(1\-S\)^\{\\lambda\}\}\+j\(1−S\)δ​Ntotγ​Gη\+e\.\\displaystyle\\quad\+\\;\\frac\{j\}\{\(1\-S\)^\{\\delta\}\\,N\_\{\\text\{tot\}\}^\{\\gamma\}\\,G^\{\\eta\}\}\\;\+\\;e\.\(10\)HereGGappears only in the denominator of the interaction term, so raisingGGshrinks the sparsity penalty and lowers the predicted loss at fixed\(Ntot,S,D\)\(N\_\{\\text\{tot\}\},S,D\)\. PlacingGGin the joint coupling rather than folding it into a single effective parameter count is what keeps the fit balanced acrossGGvalues\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/x5.png)Figure 5:Loss against total MoE parametersNtotN\_\{\\text\{tot\}\}\. In the left, one curve per expert split factorGGat fixedS=0\.95S\\\!=\\\!0\.95,D=6D\\\!=\\\!6B\. In the right, one curve per sparsity band at fixedG=4G\\\!=\\\!4,D=6D\\\!=\\\!6B\. The curves fan out rather than staying parallel \(left\), which is whyGGenters the law as its own axis instead of being absorbed into an effective parameter count, and Table[2](https://arxiv.org/html/2608.10605#S4.T2)bears this out, where theGG\-exponentη\\etais identified rather than sitting on the coefficient ridge\.Coeff\.Value95%95\\%BCa CIIdentifiabilityaa8\.4628\.462\[8\.24,8\.85\]\[8\.24,\\,8\.85\]ridgebb88\.3288\.32\[87\.8,89\.9\]\[87\.8,\\,89\.9\]ridgecc0\.18280\.1828\[0\.181,0\.184\]\[0\.181,\\,0\.184\]ridgejj0\.71230\.7123\[0\.710,0\.713\]\[0\.710,\\,0\.713\]ridgeee0\.31290\.3129\[0\.307,0\.316\]\[0\.307,\\,0\.316\]ridgeα\\alpha0\.10480\.1048\[0\.104,0\.106\]\[0\.104,\\,0\.106\]ridgeβ\\beta0\.20700\.2070\[0\.206,0\.208\]\[0\.206,\\,0\.208\]identifiedλ\\lambda0\.12490\.1249\[0\.119,0\.138\]\[0\.119,\\,0\.138\]ridgeδ\\delta0\.55670\.5567\[0\.550,0\.577\]\[0\.550,\\,0\.577\]ridgeγ\\gamma0\.17020\.1702\[0\.167,0\.174\]\[0\.167,\\,0\.174\]ridgeη\\eta\+0\.9513\\mathbf\{\+0\.9513\}\[\+0\.911,\+0\.954\]\\mathbf\{\[\+0\.911,\\,\+0\.954\]\}identifiedTable 2:Fitted coefficients of the joint scaling law in Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\)\.95%95\\%BCa CIs come from a warm\-started wild bootstrap in the deployed basin\. We work through which coefficients the data identify and which sit on the flat ridge in Appendix[B](https://arxiv.org/html/2608.10605#A2)\.##### Fitting Results\.

We fit the form in Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) on the scaling grid spanningG∈\{1,2,4,6,8\}G\\\!\\in\\\!\\\{1,2,4,6,8\\\}\. We hold out the top\-10%10\\%of runs byNactN\_\{\\text\{act\}\}within eachGG\-stratum to validate the observations we do not use in training\. We minimize the Huber objective onlog⁡ℒ\\log\\mathcal\{L\}\(δHuber=0\.1\\delta\_\{\\text\{Huber\}\}=0\.1\) by L\-BFGS\-B in log\-coefficient space\. Table[2](https://arxiv.org/html/2608.10605#S4.T2)reports the point estimates together with95%95\\%confidence intervals, and marks each coefficient as identified or ridge\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/granularity_pred_vs_real.png)Figure 6:Training loss against expert split factorG=dff/dexpertG\\\!=\\\!d\_\{\\text\{ff\}\}/d\_\{\\text\{expert\}\}at matched\(Ntot,S,D\)\(N\_\{\\text\{tot\}\},S,D\)operating points\. The predicted loss decreases monotonically withGG\. The observed runs decrease up toG=4G\\\!=\\\!4and flatten or slightly reverse atG=8G\\\!=\\\!8within run noise\. The gap between the coarsest and finest experts widens along the plotted series, tracking their rising sparsity\.A coefficient is identified when the data pin its value uniquely, and it sits on a ridge when only a combination of several coefficients is fixed\. Table[2](https://arxiv.org/html/2608.10605#S4.T2)shows thatη\\etaandβ\\betaare identified here, becauseGGandDDeach enter the law in a single place that no other coefficient can absorb\. TheGGexponentη≈0\.95\\eta\\\!\\approx\\\!0\.95is therefore our robust observation, and since it multiplies the joint sparsity and capacity term, its reading is that largerGGdiscounts the loss penalty that sparsity imposes at fixed capacity\. We do not claim a universal expert split factor law with our functional form fits, since our scaling data traverses the discrete spaceG∈\{1,2,4,6,8\}G\\\!\\in\\\!\\\{1,2,4,6,8\\\}\. We check this against the runs in Figure[6](https://arxiv.org/html/2608.10605#S4.F6), where the predictions of Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) decrease monotonically withGGat matched\(Ntot,S,D\)\(N\_\{\\text\{tot\}\},S,D\)while the observed losses decrease up toG=4G\\\!=\\\!4and flatten or slightly reverse atG=8G\\\!=\\\!8within run\-to\-run noise\. The law attributes the widening gap between the coarsest and finest experts to the rising sparsity of the larger operating points through the\(1−S\)−δ\(1\-S\)^\{\-\\delta\}factor rather than toNtotN\_\{\\text\{tot\}\}itself, since at matchedSStheNtot−γN\_\{\\text\{tot\}\}^\{\-\\gamma\}decay of the interaction term shrinks the gap as models grow\. From the systems viewpoint of this paper,η\\etaquantifies the rate at which finer\-grained experts soften the sparsity penalty at a fixed memory footprint, a benefit that the performance model in Section[5](https://arxiv.org/html/2608.10605#S5)will price against sparse MoE layer communications overhead\.

We further validate the fit quality in Figure[7](https://arxiv.org/html/2608.10605#S4.F7), which collapses a subset of all runs acrossG∈\{1,2,4,6,8\}G\\\!\\in\\\!\\\{1,2,4,6,8\\\}onto the identity line, and the joint law attains a held\-out RMSE of0\.007790\.00779in loss units on the largest10%10\\%of runs byNactN\_\{\\text\{act\}\}within eachGGstratum\. The fitted form also behaves as we expect of a well\-formed scaling law\. It is monotone\-decreasing inNtotN\_\{\\text\{tot\}\},DD, andGG, and convex inSSwith a sharp penalty asS→1S\\\!\\to\\\!1\. AsG→∞G\\\!\\to\\\!\\inftythe joint term drops out and the loss settles ontoa/Ntotα\+b/Dβ\+c​\(1−S\)−λ\+ea/N\_\{\\text\{tot\}\}^\{\\alpha\}\+b/D^\{\\beta\}\+c\\,\(1\-S\)^\{\-\\lambda\}\+e\. AsNtot,D→∞N\_\{\\text\{tot\}\},D\\\!\\to\\\!\\inftyat fixed MoE sparsitySS, the loss approachese\+c​\(1−S\)−λe\+c\\,\(1\-S\)^\{\-\\lambda\}, so the sparsity floor persists at infinite scale and the irreducible term ise∞=e\+c≈0\.50e\_\{\\infty\}\\\!=\\\!e\+c\\\!\\approx\\\!0\.50in the dense limitS→0S\\\!\\to\\\!0rather thaneealone\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_joint.png)Figure 7:Joint\-law parity plot over all runs,R2=0\.99877R^\{2\}=0\.99877\. The crosses are the held\-out validation runs, the largest10%10\\%byNactN\_\{\\text\{act\}\}within eachGGstratum, with a held\-out RMSE of0\.007790\.00779in loss units\. Marker shape and color pick out a sample of the\(G,K\)\(G,K\)pairs the sweep covers\.

## 5EstimatingMFU\(\.\)MFU\(\.\)with a Performance Model

This section specifies the design of a performance estimator that predictsMFU​\(Z,Pℓ,B\)\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)andMem​\(Z,Pℓ,B\)\\mathrm\{Mem\}\(Z,P\_\{\\ell\},B\)for MOSAIC in Section[2](https://arxiv.org/html/2608.10605#S2)\. This estimator takes only the model specification and the physical device layout as input, and produces a parallelism layout, single\-iteration running time, MFU, and memory footprint as output\. A closed\-form scaling rule cannot stand in for MFU predictions, because systems scaling is nonlinear and MFU is therefore nonlinear in the device count\. So, doubling the GPUs neither guarantees doubling the throughput nor preserves the layout that was optimal before, so we cost the forward and backward passes of a single iteration analytically against measured collective timings instead\.

### 5\.1Design Principle

The core principle behind our performance model is to minimize MFU prediction error while keeping the model simple and physically\-grounded: every term must trace to a hardware specification, a measured kernel or collective timing, or a scheduling equation, never a hand\-tuned constant\. We pair a lightweight analytical skeleton \(FLOPs, collective\-algorithm equations, and the 1F1B pipeline schedule\) with a set of carefully designed microbenchmarks run on the target hardware, so that the model consumes*measured*kernel efficiencies and collective bandwidths rather than attempting to reproduce the intrinsic complexity of runtime kernel behavior from first principles\. Two effects cannot be captured by an uncontended microbenchmark, namely CPU\-dispatch\-bound execution on small models and cross\-node straggler and overlap in large collectives\. We isolate each into a small, explicitly labeled correction rather than spreading fudge factors across every component\.

In the large, this design yields a practical calculator that guides performance engineering toward the optimal MFU through an efficient parallelism\-layout search\. In the small, it also offers a fine\-grained per\-component time breakdown of a single MoE training or inference iteration, which supports bottleneck diagnosis\.

### 5\.2Functionality

The performance model is driven by a model\-architecture specification that fully describes a training run, as shown in Appendix[E](https://arxiv.org/html/2608.10605#A5)\. It groups settings into seven blocks\. Themodelblock defines the dense transformer geometry \(depth, widths, attention heads, vocabulary, and architectural switches such as SwiGLU, rotary embeddings, and weight tying\), along with the compute precision, which selects among floating\-point formats \(bf16,fp16,fp8,fp4,tf32,fp32, and others\) so that GEMMs and attention are costed against the matching device peak\-FLOPs lane\. Themodel\.moeblock adds the mixture\-of\-experts layer \(expert count, top\-kkrouting, expert width, shared\-expert width, and the capacity factor that sizes the routed token budget and the all\-to\-all traffic\)\. Thesearchblock declares the device budget\. Theperformanceblock sets the activation\-checkpointing strategy \(none,selective,super\-selective, orfull\)\. Theoptimizerblock selects the optimizer family \(adam,adamw,muon,scion, and more\), which sets the optimizer\-state footprint and update FLOPs\. Thedatablock gives the batch and sequence shape, and thehardwareblock names the target node type\.

From this specification, the tool builds an internal model and hardware representation, then enumerates the feasible execution plansPℓ=P\_\{\\ell\}=\(PTPP\_\{\\text\{TP\}\},PEPP\_\{\\text\{EP\}\},PCPP\_\{\\text\{CP\}\},PPPP\_\{\\text\{PP\}\},PDPP\_\{\\text\{DP\}\},BmicroB\_\{\\text\{micro\}\}\) that tile the device budget and estimates the compute, memory, and communication cost of each\. The output is a per\-layout prediction of MFU and memory usage, used to rank parallelism configurations against realistic training jobs\. We give a worked example of both the specification and the returned output in Appendix[E](https://arxiv.org/html/2608.10605#A5), where we also describe the layout search that produces the ranking\.

### 5\.3Microbenchmarks

Closed\-form FLOPs divided by device peak FLOPs over\-predict throughput: real kernels reach only a fraction of peak that depends on operand shape, and memory\-bound operations do not scale with FLOPs at all\. We therefore benchmark the dominant computation and communication components on the target hardware and feed the measured timings into the model as interpolated lookup tables\. We interpolate within the measured range and refuse to extrapolate silently outside it\. The benchmarked components are:

- •Dense GEMMs: achieved utilization as a function of\(M,N,K\)\(M,N,K\)and precision, for the attention, MLP, shared\-expert, and LM\-head projections\.
- •Grouped GEMMs\(the MoE expert MLP\): the fusedGroupedLinearoperator, keyed by the number of local experts \(num\_gemms\), the per\-expert token load \(the grouped tileMM\), and the expert shape\. This is essential because the expert MLP runs many small per\-expert tiles rather than one large GEMM, where its efficiency is set by the*per\-group*tileMexpert=⌈Tseq​Bmicro​K/E⌉M\_\{\\text\{expert\}\}=\\lceil T\_\{\\text\{seq\}\}\\,B\_\{\\text\{micro\}\}\\,K/E\\rceil\(sequence lengthTseqT\_\{\\text\{seq\}\}, micro\-batchBmicroB\_\{\\text\{micro\}\}, top\-KK,EEexperts\), which is invariant to expert and data parallelism, not by the aggregate token count\. Costing it as a single dense GEMM over\-credits utilization by up to an order of magnitude in the token\-dropless regime\.
- •Scaled dot\-product attention\(SDPA\) with the Flash\-Attention backend\(zadouri2026flashattention\), forward and backward, indexed by the TP\-sharded head counts that actually execute on one GPU\.
- •MoE top\-kkroutingandelement\-wiseoperations \(silu,residual\_add\), forward and backward\.
- •Cross\-entropy over the vocabulary: the fused vocab\-parallel loss onfp32\-upcast logits\. This is far more expensive than a naive count of logit passes suggests, and it lands entirely on the last pipeline stage\.
- •NCCL collectives:all\_reduce,reduce\_scatter,all\_gather,all\_to\_all, andsend\_recv, as effective bus\-bandwidth curves swept over message size*and*group size, covering both the intra\-node \(NVSwitch\) and inter\-node \(EFA\) regimes and both contiguous \(dense\-DP\) and strided \(expert\-DP\) rank layouts\. These curves differ from the analytic link\-peak\-times\-protocol\-efficiency estimate by up to8×8\\timesin both directions and exhibit a strong size ramp that specifications alone do not predict\.
- •Per\-kernel CPU dispatch gap: the steady\-state CPU enqueue stall between kernels in an eager execution stream, which bounds throughput when kernels are small \(see below\)\.

### 5\.4Iteration\-Time Model

The predicted per\-iteration time sums a transformer\-block compute term, a pipeline\-bubble term, a vocabulary\-stage \(embedding and LM\-head/cross\-entropy\) term, the exposed data\-parallel and pipeline\-parallel collective times, and the optimizer step\. Three parts of this decomposition depart from a naive roofline and warrant explanation\.

##### Compute versus dispatch \(per\-microbatch wall\)\.

For small models with a micro\-batch size of one, the transformer\-block kernels are tiny \(e\.g\.<0\.1<0\.1ms per GEMM\) and a single microbatch launches tens of thousands of kernels\. The CPU cannot enqueue them fast enough and the GPU starves\. We therefore model the per\-microbatch wall as

tmb=max⁡\(tcmp,tdisp\),tblock=nmb​tmb,nmb=BglobalPDP​Bmicrot\_\{\\text\{mb\}\}=\\max\\\!\\big\(t\_\{\\text\{cmp\}\},\\ t\_\{\\text\{disp\}\}\\big\),\\qquad t\_\{\\text\{block\}\}=n\_\{\\text\{mb\}\}\\,t\_\{\\text\{mb\}\},\\qquad n\_\{\\text\{mb\}\}=\\frac\{B\_\{\\text\{global\}\}\}\{P\_\{\\text\{DP\}\}\\,B\_\{\\text\{micro\}\}\}\(11\)wherenmbn\_\{\\text\{mb\}\}is the number of microbatches per optimizer step,tcmpt\_\{\\text\{cmp\}\}is the summed measured GPU\-compute time of the stage’s layers, andtdisp=Lstage​κ​\(Eloc\)​tgap\+thostt\_\{\\text\{disp\}\}=L\_\{\\text\{stage\}\}\\,\\kappa\(E\_\{\\text\{loc\}\}\)\\,t\_\{\\text\{gap\}\}\+t\_\{\\text\{host\}\}is the CPU\-dispatch floor overLstageL\_\{\\text\{stage\}\}layers per stage, with a measured per\-kernel gaptgapt\_\{\\text\{gap\}\}, a fixed per\-microbatch host termthostt\_\{\\text\{host\}\}, and a kernel\-count lawκ​\(Eloc\)=κ0\+κ1​Eloc\\kappa\(E\_\{\\text\{loc\}\}\)=\\kappa\_\{0\}\+\\kappa\_\{1\}\\,E\_\{\\text\{loc\}\}that grows with the number of local expertsElocE\_\{\\text\{loc\}\}\(the surrounding framework kernels, namely permutation, copies, and bias\-adds, dominate the launch count, not the fused expert GEMM\)\. Themax\\maxmakes this term*self\-limiting*, where large\-model kernels are large enough to hide dispatch, so the term vanishes and those configurations remain compute\-bound\.

##### Pipeline bubble\.

With 1F1B scheduling, the pipeline fill and drain phases leave each stage idle for an effective fraction of the iteration\. Fornmbn\_\{\\text\{mb\}\}microbatches per pipeline,PPPP\_\{\\text\{PP\}\}pipeline stages, andPVPPP\_\{\\text\{VPP\}\}\-way interleaving444https://docs\.nvidia\.com/nemo/megatron\-bridge/0\.2\.0/parallelisms\.html\#interleaved\-pipeline\-parallel\-schedule, the bubble fraction is

fbubble=PPP−1PVPP​nmbf\_\{\\text\{bubble\}\}=\\frac\{P\_\{\\text\{PP\}\}\-1\}\{P\_\{\\text\{VPP\}\}\\,n\_\{\\text\{mb\}\}\}\(12\)so the pipeline bubble time istbubble=fbubble​tblockt\_\{\\text\{bubble\}\}=f\_\{\\text\{bubble\}\}\\,t\_\{\\text\{block\}\}\. Each rank computesnmbn\_\{\\text\{mb\}\}microbatches of work and idles\(PPP−1\)/PVPP\(P\_\{\\text\{PP\}\}\-1\)/P\_\{\\text\{VPP\}\}microbatch\-slots during warmup and cooldown\. That ramp is the entire bubble\. This form replaces an earlier weighted\-average expression that added a separate fill/drain term, charging each of the2​\(PPP−1\)2\(P\_\{\\text\{PP\}\}\-1\)warmup/cooldown microbatches a\(PPP−1\)/PPP\(P\_\{\\text\{PP\}\}\-1\)/P\_\{\\text\{PP\}\}stage bubble on top of the\(PPP−1\)/nmb\(P\_\{\\text\{PP\}\}\-1\)/n\_\{\\text\{mb\}\}steady\-state term, and clamped the result at\(PPP−1\)/PPP\(P\_\{\\text\{PP\}\}\-1\)/P\_\{\\text\{PP\}\}\. That expression double\-counted the fill/drain phase, and its clamp under\-counted the small\-batch regime it was meant to bound\.

##### Vocabulary\-stage imbalance\.

The token embedding is placed on the first pipeline stage and the LM head plus cross\-entropy is placed on the last, in addition to those stages’ transformer layers\. Because 1F1B is paced by the slowest stage, the surplustlastt\_\{\\text\{last\}\}on the last stage \(the measured LM\-head GEMM forward and backward plus the measured cross\-entropy\) is exposed once per microbatch, and it further lengthens the fill/drain ramp\. The stage\-imbalance contribution is

tvocab=\(nmb−1\)​max⁡\(tfirst,tlast\)\+tfirst\+tlastt\_\{\\text\{vocab\}\}=\(n\_\{\\text\{mb\}\}\-1\)\\,\\max\(t\_\{\\text\{first\}\},t\_\{\\text\{last\}\}\)\+t\_\{\\text\{first\}\}\+t\_\{\\text\{last\}\}\(13\)with an addedtlast​nmb​fbubblet\_\{\\text\{last\}\}\\,n\_\{\\text\{mb\}\}\\,f\_\{\\text\{bubble\}\}term for the ramp\. Heretfirstt\_\{\\text\{first\}\}is the \(bandwidth\-bound\) embedding gather/scatter\. This term is essential for deep\-pipeline MoE configurations, where the vocabulary\-parallel LM head is a large fraction of the critical path and was previously unpriced\.

The exposed data\-parallel and pipeline\-parallel collective times are taken directly from the measured bandwidth curves of Section[5\.3](https://arxiv.org/html/2608.10605#S5.SS3), with the DP grad reduce\-scatter and param all\-gather overlapped against compute where windows allow\.

### 5\.5Calibration

Because the compute and wire\-time terms above are read from measured lookup tables, most of the model requires no correction factor at all\. What remains, and cannot be recovered from an uncontended microbenchmark, is a single*runtime*effect, namely cross\-node straggler and overlap\. When an expert\-parallel or pipeline\-parallel group spans multiple nodes, many collective rings run concurrently and are paced by the slowest rank \(traces show per\-call inter\-node latencies with a tight median but tails two to three orders of magnitude larger\), so the realized cost exceeds the uncontended microbenchmark\. When a deep pipeline fits within a single node, conversely, the heavy last stage overlaps other stages over NVLink better than a straggler\-free 1F1B accounting assumes, so the realized step is*faster*than predicted\.

We isolate this into one explicitly labeled correction, the model’s only fitted term, which multiplies the analytical iteration time,

titercal=χsys​titeranalytict^\{\\text\{cal\}\}\_\{\\text\{iter\}\}=\\chi\_\{\\text\{sys\}\}\\,t^\{\\text\{analytic\}\}\_\{\\text\{iter\}\}\(14\)where the systems calibration multiplier is

χsys=1\+cA2A​tA2Atiter​1​\[NnodeEP\>1\]\+cPP​\(NnodePP−1\)\+covl​1​\[PPP≥Ngpu/node,NnodePP=1\]\.\\begin\{split\}\\chi\_\{\\text\{sys\}\}&=1\+c\_\{\\text\{A2A\}\}\\,\\frac\{t\_\{\\text\{A2A\}\}\}\{t\_\{\\text\{iter\}\}\}\\,\\mathbb\{1\}\\\!\\left\[N^\{\\text\{EP\}\}\_\{\\text\{node\}\}\>1\\right\]\+c\_\{\\text\{PP\}\}\\,\\big\(N^\{\\text\{PP\}\}\_\{\\text\{node\}\}\-1\\big\)\\\\ &\\quad\+c\_\{\\text\{ovl\}\}\\,\\mathbb\{1\}\\\!\\left\[P\_\{\\text\{PP\}\}\\geq N\_\{\\text\{gpu/node\}\},\\ N^\{\\text\{PP\}\}\_\{\\text\{node\}\}=1\\right\]\.\\end\{split\}\(15\)HereNnodeEP=⌈PEP​PTP/Ngpu/node⌉N^\{\\text\{EP\}\}\_\{\\text\{node\}\}=\\lceil P\_\{\\text\{EP\}\}\\,P\_\{\\text\{TP\}\}/N\_\{\\text\{gpu/node\}\}\\rceilandNnodePP=⌈PPP/Ngpu/node⌉N^\{\\text\{PP\}\}\_\{\\text\{node\}\}=\\lceil P\_\{\\text\{PP\}\}/N\_\{\\text\{gpu/node\}\}\\rceilcount how many nodes ofNgpu/nodeN\_\{\\text\{gpu/node\}\}GPUs each the expert\-parallel and pipeline\-parallel groups span, andtA2At\_\{\\text\{A2A\}\}is the modeled per\-step expert all\-to\-all wire time\. Each driver is tied to a physical quantity, where \(i\) the inter\-node A2A straggler scales the penalty by how large a fraction of the step the expert A2A already is, and applies only when the EP group crosses a node boundary\. \(ii\) The deep\-pipeline term scales with the number of inter\-node hops the pipeline crosses\. \(iii\) The intra\-node overlap*credit*\(covl<0c\_\{\\text\{ovl\}\}<0\) applies only to a deep pipeline confined to one node\. Single\-node, shallow\-pipeline configurations incur no correction \(χsys=1\\chi\_\{\\text\{sys\}\}=1\)\.

Together, the microbenchmarks, the iteration\-time model, and this single calibration term constitute the estimator we release as ScalePlan, whose layout search supplies theMFU\\mathrm\{MFU\}andMem\\mathrm\{Mem\}predictors that MOSAIC consumes\. We validate these predictions against measured MFU on real training sweeps in Section[7\.2](https://arxiv.org/html/2608.10605#S7.SS2)\.

## 6Optimal MoE Sparsity Scaling

In Section[2](https://arxiv.org/html/2608.10605#S2), we defined deliverable model FLOPsCdeliverableC\_\{\\text\{deliverable\}\}as the systems\-aware constraint used by MOSAIC\. We now develop the motivation sketched in Figure[1](https://arxiv.org/html/2608.10605#S1.F1)in more detail for sparse MoEs\. We bring together the two pieces, the sparse MoE scaling laws and the hardware\-deliverable model FLOPs, to explain why the deliverable model FLOPsCdeliverableC\_\{\\text\{deliverable\}\}of Equation \([2](https://arxiv.org/html/2608.10605#S2.E2)\) is the budget that makes sense for sparse MoEs before we turn to results\. A practical use of scaling laws such as ours is to pick the scaling dimensions that minimize predicted loss under a model FLOPs budget,Cmodel=6​Nact​DC\_\{\\text\{model\}\}\\\!=\\\!6\\,N\_\{\\text\{act\}\}D\. Simply said, we want the\(Ntot,S,D,G\)\(N\_\{\\text\{tot\}\},S,D,G\)that minimizes our law under the budget\. Conventionally, solutions of this kind proceed in two stages\. In the first stage, we solve an optimization for the model FLOPs optimal\(Ntot⋆,S⋆,D⋆,G⋆\)\(N\_\{\\text\{tot\}\}^\{\\star\},S^\{\\star\},D^\{\\star\},G^\{\\star\}\), treating these as continuous variables\. In the second stage, we map that optimum to a discrete, deployable model configuration, sinceSSandGGare not free knobs one can set directly\.SSis a derived ratio, andGGmoves only by changing the discrete geometry\(E,K,Llayers,d\)\(E,K,L\_\{\\text\{layers\}\},d\)of Section[3](https://arxiv.org/html/2608.10605#S3)\. To that end, Sections[6\.1](https://arxiv.org/html/2608.10605#S6.SS1)and[6\.2](https://arxiv.org/html/2608.10605#S6.SS2)discuss two limitations of model\-FLOPs based optimality for sparse MoE scaling that make such a two\-stage optimization the wrong choice\. The next section shows that the sparse MoE dimensions in our scaling law admit no interior optimum under a model\-FLOPs budget, which forces the budget to be recast as an architecture\-dependent deliverable model\-FLOPs budget\.

### 6\.1Model\-FLOPs Optimal Sparsity Is Boundary\-Seeking

Conditioning on the scaling law, the conventional scaling studies seek theoptimal sparsitySSthat minimizes loss at a fixed budgetCmodelC\_\{\\text\{model\}\}\. Concretely, withNact=\(1−S\)​NtotN\_\{\\text\{act\}\}=\(1\-S\)\\,N\_\{\\text\{tot\}\}and the budgetCmodel=6​Nact​DC\_\{\\text\{model\}\}=6\\,N\_\{\\text\{act\}\}\\,D\(at fixedGG\), the model FLOPs only optimization problem reads

minNtot,S,D⁡ℒ​\(Ntot,S,D,G\)s\.t\.6​\(1−S\)​Ntot​D=Cmodel,Ntot,D\>0,S∈\[0,1\)\.\\min\_\{N\_\{\\text\{tot\}\},S,D\}\\;\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D,G\)\\quad\\text\{s\.t\.\}\\quad 6\\,\(1\-S\)\\,N\_\{\\text\{tot\}\}\\,D\\;=\\;C\_\{\\text\{model\}\},\\quad N\_\{\\text\{tot\}\},D\>0,\\quad S\\in\[0,1\)\.\(16\)Solving Equation \([16](https://arxiv.org/html/2608.10605#S6.E16)\) with the joint law of Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) atCmodel∈\{3×1019,6×1019,1020,3×1020,6×1020\}C\_\{\\text\{model\}\}\\in\\\{3\\times 10^\{19\},\\,6\\times 10^\{19\},\\,10^\{20\},\\,3\\times 10^\{20\},\\,6\\times 10^\{20\}\\\}FLOPs recovers the three trends that theG=4G\\\!=\\\!4fit of Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\) produces in Figure[4](https://arxiv.org/html/2608.10605#S4.F4)\. The important observation in the plots is thatNtot⋆​\(S\)N\_\{\\text\{tot\}\}^\{\\star\}\(S\)increases monotonically withSS,Nact⋆​\(S\)N\_\{\\text\{act\}\}^\{\\star\}\(S\)decreases, andℒ⋆​\(S\)\\mathcal\{L\}^\{\\star\}\(S\)decreases monotonically at everyCmodelC\_\{\\text\{model\}\}\. Becauseℒ⋆​\(S\)\\mathcal\{L\}^\{\\star\}\(S\)decreases monotonically over the calibrated sparsity range, the interior condition∂ℒ⋆/∂S=0\\partial\\mathcal\{L\}^\{\\star\}/\\partial S=0is not satisfied within this region\. The optimal sparsity under theCmodelC\_\{\\mathrm\{model\}\}budget, therefore lies atS=SmaxS=S\_\{\\max\}, the upper boundary of our data support\. Thus, within this range where the scaling law is calibrated, model\-FLOPs only optimization reduces to “makeSSas large as the supported architecture space allows\.” The left column of Figure[8](https://arxiv.org/html/2608.10605#S6.F8)shows the same monotone descent on both knobs, inSSat fixedGGand inGGat fixedSS\. We do not claim that this monotonic trend continues arbitrarily close toS=1S=1, where the asymptotic sparsity\-penalty terms in the fitted law eventually dominate\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/x6.png)Figure 8:Opposing views of model scaling and systems cost\.The top row varies MoE sparsitySSholdingG=4G\\\!=\\\!4\. The bottom row varies the expert split factorGGholdingS=0\.958S\\\!=\\\!0\.958\. The left column gives the optimal loss under a fixed model\-FLOPs budget\. Each marker there is labeled with the optimal total parametersNtotN\_\{\\text\{tot\}\}under that budget\. The right column shows the wall\-clock time each configuration needs per delivered model FLOP on512512devices \(64 p6 nodes\), relative to the training time of the lowest sparsity configuration\. TheCmodelC\_\{\\text\{model\}\}optimal loss configuration pays1\.70×1\.70\\timesper FLOP in the top row and2\.66×2\.66\\timesin the bottom\. The starred winners sit at opposite ends, so theCmodelC\_\{\\text\{model\}\}optimal loss configuration is the systems\-slowest\.
### 6\.2The Boundary Optimum Is Not Necessarily the Cluster Optimum

In the preceding section, the model\-FLOPs optimum proved boundary\-seeking in sparsity, returning the largestSSthe architecture space supports rather than an interior choice\. Additionally and more importantly, practitioners do not buy model FLOPs, they buy GPU\-hours on a fixed cluster over a fixed wall\-clock window, and what they care about at the end of training is the training loss reached\. TheCdeliverableC\_\{\\text\{deliverable\}\}the cluster can execute thus becomes the metric that matters once we account for the systems cost of hosting a model across a fleet of GPUs\(shoeybi2019megatron\), which is what we show in Figure[1](https://arxiv.org/html/2608.10605#S1.F1)\. In particular, sharding a model through various parallel layouts incurs costs from expert\-parallel all\-to\-all, pipeline bubbles, optimizer sharding, and host\-side stalls\(singh2023hybrid;fernandez2024hardware\)\. What follows is that a model FLOPs budget leaves out the systems\-side constraints that together with device peak FLOPs provideCdeliverableC\_\{\\text\{deliverable\}\}\.

To make the motivation concrete, we take as the systems metric the wall\-clock time a configuration needs per delivered model FLOP, where a lower value means more of the purchased hardware is doing useful work and the same loss is reached in fewer GPU\-hours\. The top row of Figure[8](https://arxiv.org/html/2608.10605#S6.F8)shows how loss, sparsity, and execution cost trade off on the way to the deployable optimum\. Its left panel plots the optimal lossℒ⋆\\mathcal\{L\}^\{\\star\}reached at each sparsity level at fixedCmodel=1022C\_\{\\text\{model\}\}\\\!=\\\!10^\{22\}model FLOPs, after minimizing over active parameters\. It drops monotonically as sparsity grows, so under a model FLOPs budget the highest\-sparsity configuration ranks first and the lowest\-sparsity one ranks last\. The right panel prices the same grid at world\-size512512\(number of GPU devices\) with the performance model of Section[5](https://arxiv.org/html/2608.10605#S5), optimized over the parallelism layout, and the wall\-clock per model FLOP rises monotonically with sparsity to1\.70×1\.70\\timesthe densest configuration\. This opposing ordering is precisely why the architecture\-dependent deliverable model\-FLOPs constraint introduced in Section[2](https://arxiv.org/html/2608.10605#S2)is required for sparse MoEs\.

An important point to note here is that, unlike the sparsity scaling laws behind the model design report of Kimi K2\(team2025kimi\), which fix the active parameters, or the recommendation ofabnar2025parametersthat at a fixed total parameter count a larger training budget lowers the optimal sparsity, we hold neitherNactN\_\{\\text\{act\}\}norNtotN\_\{\\text\{tot\}\}fixed in Figure[8](https://arxiv.org/html/2608.10605#S6.F8)\. We instead propose a shift toward hardware constraints of Section[2](https://arxiv.org/html/2608.10605#S2)to resolve these optimality questions\.

## 7Experiments and Results

In this section, we demonstrate how we utilize MOSAIC towards the goal of searching for optimal model configuration and the parallel execution layout jointly under a hardware budget\. We utilize the aforementioned components of the scaling laws and the performance model towards this goal\. We detail how this search solves the program of Section[2\.2](https://arxiv.org/html/2608.10605#S2.SS2)in Appendix[D](https://arxiv.org/html/2608.10605#A4)exactly through structure\-exploiting enumeration\. We note that generic MINLP solvers are poorly matched to it, since the MFU and memory constraints are measured lookup surfaces rather than symbolic functions\. To surface how a hardware\-aware optimum differs from a purely model FLOPs optimal one, we first discuss the need for a geometry scaling ladder in Section[7\.1](https://arxiv.org/html/2608.10605#S7.SS1)\. SparsitySSand the expert split factorGGare derived from the discrete geometry rather than set directly, so the search needs a closed set of realizable architectures, which the ladder supplies and the performance model then profiles for systems cost\. Our scaling law sweep runs also share the same ladder as every experiment reported here, so one family of geometries both fits the law and is searched by the framework for validation\. Every experiment also uses the joint lawℒ​\(Ntot,S,D,G\)\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D,G\)of Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\), not the per\-GGfits of Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\)\.

After presenting the ladder, we validate the performance model against measured MFU in Section[7\.2](https://arxiv.org/html/2608.10605#S7.SS2), since the whole framework rests on its predictions\. The co\-design results follow in three parts, where the first two search the full feasible grid with the framework alone and the third launches its ranked predictions as real training runs\. These comprise \(1\) the loss\-optimal configuration and the loss\-versus\-model\-FLOPs spread over feasible configurations under a fixedCpeakC\_\{\\rm peak\}envelope in Section[7\.3](https://arxiv.org/html/2608.10605#S7.SS3), \(2\) the optimal MoE\-layer geometry in sparsity and expert split factorGGas we vary the cluster size \(nodes\) and training window \(days\) in Section[7\.4](https://arxiv.org/html/2608.10605#S7.SS4), and \(3\) staged real training runs that test the predicted rank ordering against measured throughput in Section[7\.5](https://arxiv.org/html/2608.10605#S7.SS5)\. A fitted loss surface can be validated against the values it predicts or against the configuration ordering it induces\. This ordering is what a practitioner acts on when choosing model configurations, which is why we end the results section with real runs selected under a common hardware envelope rather than with extrapolations of the loss surface alone\.

All our validation and training runs use AWS p6\-b200\.48xlarge nodes555https://aws\.amazon\.com/ec2/instance\-types/p6/, each with88NVIDIA B200 GPUs\. We provide the details of the Megatron and Torch\-based framework we stage our runs on in Appendix[F\.3](https://arxiv.org/html/2608.10605#A6.SS3), covering which Megatron\-Core MoE features we enable, the loop\-level effects that make realized MFU trail the analytical prediction, the host\-side data\-pipeline costs, and the hardware and interconnect assumptions that tie the recovered optimum to a specific platform\.

### 7\.1Geometry Scaling Ladder

Since we want to show how the hardware\-aware optimum shifts as the deliverable model FLOPs change with scale, we first need a way to climb the model\-scale axis\. Growing a model means growing depth, width, and FFN width together, so that axis is well defined only once we fix how those dimensions co\-scale\. A geometry scaling ladder is that choice, and it reduces the backbone to a single seed\. This matters for MOSAIC, which selects a discrete geometry in a single step rather than through the traditional two\-stage model\-FLOPs allocation, so the geometry must be pinned down directly\. We parametrize the ladder in the following way: the layer count and head count both followLlayers=m​qL\_\{\\text\{layers\}\}=m\\,qandnhead=m​qn\_\{\\text\{head\}\}=m\\,qwith the multipliermmfixed to 4, so the residual width isd=nhead​dheadd=n\_\{\\text\{head\}\}\\,d\_\{\\text\{head\}\}atdhead=128d\_\{\\text\{head\}\}=128, and the reference dense FFN width isdff=2\.5​dd\_\{\\text\{ff\}\}=2\.5\\,d\. A single seedqqthus fixes depth \(LlayersL\_\{\\text\{layers\}\}\), width \(dd\), and FFN width together at a constant depth\-to\-width ratio, so the ladder grows the backbone along one coherent axis while the MoE axes\(E,K,G\)\(E,K,G\), and through them the sparsitySS, are searched on top of it\. We fix the co\-scaling functions of our ladder this way and vary the geometry seedqq, and we leave other variants of co\-scaling geometry as future work\.

With the ladder in place, the framework searches over the seedqq, the expert countEE, the top\-KK, the split factorGG, and the data budgetDD, returning the optimum, and the sparsitySSand the parameter countsNtotN\_\{\\text\{tot\}\}andNactN\_\{\\text\{act\}\}are derived from the selected geometry rather than searched directly\. We show how this is realized in the following sections\. We list the model scales we sweep in Table[9](https://arxiv.org/html/2608.10605#A6.T9)of Appendix[F](https://arxiv.org/html/2608.10605#A6)\.

### 7\.2Performance Model Validation

Our validation spans three sweeps on these p6 nodes, over theq=4q\\\!=\\\!4,q=8q\\\!=\\\!8, andq=12q\\\!=\\\!12rungs of Table[9](https://arxiv.org/html/2608.10605#A6.T9)\. Theq=4q\\\!=\\\!4model \(0\.70\.7BNactN\_\{\\text\{act\}\},1616BNtotN\_\{\\text\{tot\}\}\) runs on 4 nodes \(3232GPUs\), theq=8q\\\!=\\\!8model \(5\.45\.4BNactN\_\{\\text\{act\}\},127127BNtotN\_\{\\text\{tot\}\}\) on 32 nodes \(256256GPUs\) and 64 nodes \(512512GPUs\), and theq=12q\\\!=\\\!12model \(1818BNactN\_\{\\text\{act\}\},434434BNtotN\_\{\\text\{tot\}\}\) on 64 nodes \(512512GPUs\)\. Each data point pairs one parallelism configuration, namely expert \(ep\), pipeline \(pp\), data \(dp\), tensor \(tp\), context \(cp\) parallel degree and micro batch size \(mbs\) with its measured MFU and its predicted value\. The MFU prediction results are shown in Figure[9](https://arxiv.org/html/2608.10605#S7.F9)\. We vary the global batch size and restrictt​p=1tp=1andc​p=1cp=1for efficiency\. We release the estimator as ScalePlan666[https://github\.com/dmlc/ScalePlan](https://github.com/dmlc/ScalePlan), which packages the analytical cost model, the calibration data it is fit against, and the parallelism search that realizes the inner maximization over𝒫feas\\mathcal\{P\}\_\{\\text\{feas\}\}of Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\), so the predictions validated here and every estimator number reported later can be regenerated from a single model and cluster specification\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/700m_p6-b200_32dev_gbs256_normalized.png)
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/18b_p6-b200_512dev_gbs512_normalized.png)
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/5p3b_p6-b200_256dev_gbs2048_normalized.png)
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/5p3b_p6-b200_512dev_gbs2048_normalized.png)

Figure 9:Predicted versus Real MFU per parallelism configuration, each panel normalized to its own maximum and sorted by real MFU\.#### 7\.2\.1MFU Prediction

We find that the predicted MFU values are accurate in absolute terms, not just in ordering, across all three sweeps\. The mean absolute percentage error is under15%15\\%for every model \(9\.6%9\.6\\%/9\.0%9\.0\\%/13\.2%13\.2\\%for 700M/5\.4B/18B\)\. The estimator also captures the top region of the configuration space reliably\. Ranking configurations by predicted MFU and comparing against the measured ranking, Top\-5 overlap is100%100\\%for every 5\.4B and 18B sweep and80%80\\%for the 700M sweep\. Top\-10 overlap is100%100\\%for all 5\.4B and 18B sweeps and90%90\\%for 700M\. The single 700M miss is a swap among near\-tied low\-MFU configurations, so the model reliably surfaces the high\-throughput region even where it mis\-orders adjacent points\.

#### 7\.2\.2Implications for Sparse\-MoE Parallelism Design

The MFU sweeps expose a consistent structure in how a sparse MoE responds to its parallelism axes, which we develop into design rules in Appendix[F\.2](https://arxiv.org/html/2608.10605#A6.SS2)\. The dominant lever is not any single degree but the pipeline fill ratiorfill=nmb/PPPr\_\{\\text\{fill\}\}=n\_\{\\text\{mb\}\}/P\_\{\\text\{PP\}\}, which explains almost all of the up to3\.1×3\.1\\timeslayout spread in measured MFU, with throughput collapsing asrfillr\_\{\\text\{fill\}\}falls toward11and saturating forrfill≥4r\_\{\\text\{fill\}\}\\geq 4\. The expert all\-to\-all belongs inside the scale\-up domain, since expert parallelism helps up to the node boundary and hurts past it, and the FFN\-parallelism budget is better spent on expert than expert tensor parallelism\. In practice we fixPTP=1P\_\{\\text\{TP\}\}\{=\}1, setPEPP\_\{\\text\{EP\}\}to the scale\-up domain, and choose the\(PPP,Bmicro\)\(P\_\{\\text\{PP\}\},B\_\{\\text\{micro\}\}\)pair that keepsrfill≥4r\_\{\\text\{fill\}\}\\geq 4, stated as the MFU\-optimal layout for 700M–18B models on NVL8 under no memory pressure rather than as a hardware\-independent law\.

### 7\.3Optimal Allocation UnderCpeakC\_\{\\text\{peak\}\}

In this section, we first show the configuration MOSAIC finds optimal under a single fixed GPU\-hour budget, and then extend it to the loss\-versus\-deliverable\-model\-FLOPs spread over all feasible configurations\. The inputs to MOSAIC are the training window \(wall\-clock days\), the device type and its per\-device peak FLOPs and memory, and the device count, which together set the GPU\-hour budget, alongside the global and sequence batch\.

InputsHardware32 p6 nodes \(Ndev=256N\_\{\\text\{dev\}\}=256devices\)Per\-device peakFpeakF\_\{\\text\{peak\}\}2\.252\.25PFLOPs \(bf16/dense\)Wall\-clockTtrainT\_\{\\text\{train\}\}20 daysPer\-device memory capMdevM\_\{\\text\{dev\}\}179179GBGoodputηgood\\eta\_\{\\text\{good\}\}1\.01\.0Raw hardware ceilingCpeakC\_\{\\text\{peak\}\}9\.95×10239\.95\\times 10^\{23\}FLOPsGlobal / sequence batchBglobal=1024B\_\{\\text\{global\}\}=1024,Tseq=2048T\_\{\\text\{seq\}\}=2048Loss lawjointℒ​\(Ntot,S,D,G\)\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D,G\)from Eq\. \([10](https://arxiv.org/html/2608.10605#S4.E10)\)Search space \(output\)Outer architecture gridq,E,K,Gq,E,K,GInner layout searchoverPTPP\_\{\\text\{TP\}\},PCPP\_\{\\text\{CP\}\},PEP,PPP,PDPP\_\{\\text\{EP\}\},P\_\{\\text\{PP\}\},P\_\{\\text\{DP\}\},Bmicro∈\{1,2,4\}B\_\{\\text\{micro\}\}\\\!\\in\\\!\\\{1,2,4\\\}, and activation\-checkpointing∈\{none, selective, super\-selective, full\}\\in\\\!\\\{\\text\{none, selective, super\-selective, full\}\\\}Table 3:MOSAIC search envelope\. The outer grid over\(q,E,K,G\)\(q,E,K,G\)is the discrete realization ofZZ, and the inner parallelism search realizes the embedded maximization over𝒫feas\\mathcal\{P\}\_\{\\text\{feas\}\}in Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\)\.We target a deployment simulation with 32 p6 nodes \(Ndev=256N\_\{\\text\{dev\}\}=256devices\) andTtrain=20​daysT\_\{\\text\{train\}\}=20\\,\\text\{days\}\. These inputs are listed in Table[3](https://arxiv.org/html/2608.10605#S7.T3)\. The model FLOPs the returned optimum realizes under this ceiling lie about150×150\\timesbeyond our largest fit run, an extrapolation the deployed law supports for the configuration ordering\. Combining the scaling ladder of Table[9](https://arxiv.org/html/2608.10605#A6.T9)with the sparse MoE axes, we restrict the free discrete search dimensions to

q\\displaystyle q∈\{8,10,12,14,16,18\},\\displaystyle\\in\\\{8,10,12,14,16,18\\\},K\\displaystyle K∈\{1,2,4,6,8\},\\displaystyle\\in\\\{1,2,4,6,8\\\},E\\displaystyle E∈\{4,6,8,16,24,32,48,96,128,192,256,384,512\},\\displaystyle\\in\\\{4,6,8,16,24,32,48,96,128,192,256,384,512\\\},G\\displaystyle G∈\{1,2,4,8\}\.\\displaystyle\\in\\\{1,2,4,8\\\}\.We add a practical geometry guard to prune the grid by settingS≥0\.7S\\geq 0\.7, since the scaling law data we collect to fit the loss law lies predominantly in theS≥0\.7S\\geq 0\.7region, so this keeps the search inside the sparsity range the law is calibrated on rather than extrapolating below\. The experimental envelope for this section can be found in Table[3](https://arxiv.org/html/2608.10605#S7.T3)\. We fix the token budgetDDper configuration by sweeping the tokens\-per\-parameter ratioTPP=D/Nact\\mathrm\{TPP\}=D/N\_\{\\text\{act\}\}over the gridTPP∈\[40,150\]\\mathrm\{TPP\}\\in\[40,150\]\. Our training data spans TPP up to7070\(Table[6](https://arxiv.org/html/2608.10605#A3.T6)\), so the cap sits about twice beyond the calibrated range, and we keep only\(Z,TPP\)\(Z,\\mathrm\{TPP\}\)pairs whoseCmodelC\_\{\\text\{model\}\}fits the deliverable model FLOPsCdeliverableC\_\{\\text\{deliverable\}\}at that layout’s MFU together with the memory and parallelism constraints\.

#### 7\.3\.1Results

MOSAIC returns a single loss optimal configuration, reported in Table[4](https://arxiv.org/html/2608.10605#S7.T4)\. The constrained optimum is aq=12,Llayers=48,E=96,K=2,G=4q\\\!=\\\!12,L\_\{\\text\{layers\}\}\\\!=\\\!48,E\\\!=\\\!96,K\\\!=\\\!2,G\\\!=\\\!4model withNact=14\.5N\_\{\\text\{act\}\}\\\!=\\\!14\.5B at sparsityS=0\.956S\\\!=\\\!0\.956, running at12\.4%12\.4\\%predicted MFU and lossℒ=1\.3898\\mathcal\{L\}\\\!=\\\!1\.3898\. Its selected layout isPPP=8P\_\{\\text\{PP\}\}\{=\}8,PEP=8P\_\{\\text\{EP\}\}\{=\}8,PCP=1P\_\{\\text\{CP\}\}\{=\}1,PTP=1P\_\{\\text\{TP\}\}\{=\}1with micro\-batchBmicro=4B\_\{\\text\{micro\}\}\{=\}4\.

LlayersL\_\{\\text\{layers\}\}EEKKGGNactN\_\{\\text\{act\}\}\(B\)SSTPPDD\(T\)MFU \(%\)Lossppepmbcptp48962414\.50\.956981\.4212\.381\.389888411

Table 4:The loss optimal configuration MOSAIC returns under the envelope of Table[3](https://arxiv.org/html/2608.10605#S7.T3)\. The MFU column is the predicted peak MFU of the best feasible execution plan, not a measured value\.Although Table[4](https://arxiv.org/html/2608.10605#S7.T4)reports only the loss optimal configuration, we now show all feasible\(Z,D\)\(Z,D\)configurations under the same envelope in Figure[10](https://arxiv.org/html/2608.10605#S7.F10)\. WhileDDis selected on a capped TPP grid, the cap is non\-binding for most configurations near the low\-loss frontier, whose selected token counts consume nearly the full 20\-day wall\-clock\. Near the low\-loss frontier, we observe that the geometry lands farther right primarily by sustaining a higher MFU, which converts more of the fixed hardware ceiling into useful model FLOPs\. The fixed cluster and training window emit the same raw hardware ceiling of9\.95×10239\.95\\times 10^\{23\}FLOPs for every candidate, so the quantity that varies is how much of that ceiling each geometry converts into model FLOPs through its own MFU\. Which is what we show in Figure[10](https://arxiv.org/html/2608.10605#S7.F10)\. It plots the loss against the model FLOPsCmodelC\_\{\\text\{model\}\}each candidate realizes, color\-coded byGG, with the gold star marking the loss optimal configuration\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/repro_fig_pareto.png)Figure 10:Loss versus model FLOPsCmodelC\_\{\\text\{model\}\}over all feasible\(Z,D\)\(Z,D\)configurations under the3232\-node,2020\-day envelope, color\-coded byGG\. Every candidate draws on the same raw hardware ceiling of9\.95×10239\.95\\times 10^\{23\}hardware FLOPs\. The gold star is the optimal configuration with MOSAIC, and the candidate that realizes the most model FLOPs is not the one that reaches the lowest loss\.Because the raw hardware ceiling is the same for every candidate, a geometry lands further right in Figure[10](https://arxiv.org/html/2608.10605#S7.F10)only by sustaining a higher MFU, which turns the fixed2020days of wall\-clock into more useful model FLOPs\. The expert split factorGGorders that spread\. TheG=8G\\\!=\\\!8geometries sit at the lowCmodelC\_\{\\text\{model\}\}end at a median MFU near8%8\\%, and the coarserG=2G\\\!=\\\!2geometries reach the right end near1313–20%20\\%, so a higherGGdelivers less useful compute out of the same ceiling because its expert dispatch and combine communication leaves fewer FLOPs for the model, a penalty tied to the Megatron\-Core stack our performance model mirrors \(Appendix[F\.3](https://arxiv.org/html/2608.10605#A6.SS3)\)\. Model FLOPs alone would read this as the finest experts being cheapest, but under the fixed hardware ceiling, they are the ones the cluster feeds the least\.

The consequence is that the run with the most model FLOPs is not the optimal as shown in Figure[10](https://arxiv.org/html/2608.10605#S7.F10)\. The candidate that converts the shared hardware ceilingCpeakC\_\{\\text\{peak\}\}into the most model FLOPs is a denserG=2,E=32G\\\!=\\\!2,E\\\!=\\\!32geometry at1\.94×10231\.94\\times 10^\{23\}FLOPs, and it still lands at a higher loss than the optimal configuration\. The loss optimal configuration is aG=4,E=96,K=2G\\\!=\\\!4,E\\\!=\\\!96,K\\\!=\\\!2model atNact=14\.5N\_\{\\text\{act\}\}\\\!=\\\!14\.5B andS=0\.956S\\\!=\\\!0\.956that reachesℒ=1\.3898\\mathcal\{L\}\\\!=\\\!1\.3898on1\.23×10231\.23\\times 10^\{23\}model FLOPs, roughly36%36\\%fewer than the rightmost candidate spends\. Ranking these candidates by the model FLOPs they realize therefore misorders them, and the split factor and sparsity that win are set only once the shared deliverable model\-FLOPs budget is imposed\. This is the optimum under deliverable rather than model FLOPs\.

### 7\.4OptimalSSandGG

The frontier in the previous section fixed a single cluster envelope\. We now ask how the optimal architecture moves as the deliverable model\-FLOPs budget grows, specifically whether the MoE configurations get sparser and finer, or denser and coarser under such a constraint\. To that end, we run two budget sweeps that share a common reference point\. The first fixes the cluster at3232p6 nodes \(256256GPUs\) and varies the wall\-clock window over\{10,20,30,40\}\\\{10,20,30,40\\\}days\. The second fixes the window at3030days and varies the cluster over\{16,32,64,128\}\\\{16,32,64,128\\\}p6 nodes, re\-running the performance model at each node count\. In both sweeps, we run MOSAIC subject to these constraints, and for every budget we record the architecture that minimizes predicted loss\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/repro_fig_days.png)\(a\)3232p6 nodes\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/repro_fig_nodes.png)\(b\)3030days\.

Figure 11:Optimal predicted loss vs\. the budget under the two sweeps\. Markers are colored by the optimal expert split factorG⋆G^\{\\star\}and annotated withS⋆S^\{\\star\}andNtot⋆N\_\{\\text\{tot\}\}^\{\\star\}\. Dashed lines denote best perGGconfiguration\. The solid line is the overall optimum\. The optimal sparsityS⋆S^\{\\star\}does not move monotonically with the budget in either sweep, staying on a wide plateau of roughly0\.9150\.915to0\.9630\.963\.We plot the results from running MOSAIC in Figure[11](https://arxiv.org/html/2608.10605#S7.F11)\. We observe thatG⋆G^\{\\star\}is44across most budgets, withG=8G\\\!=\\\!8appearing only at the longest window andG=2G\\\!=\\\!2at the largest node budget, and the exactG⋆G^\{\\star\}marker is often a near\-tie\. Looking at the span of sparsity, we find thatS⋆S^\{\\star\}sits on a wide plateau \(roughly0\.9150\.915–0\.9630\.963\) and does not sharpen with budget, confirming that there is no monotone sparsity trend here unlike the model FLOPs optimality we find in Figure[4](https://arxiv.org/html/2608.10605#S4.F4)\. Finally,Nact⋆N\_\{\\text\{act\}\}^\{\\star\}moves slowly with budget, staying near1414–1515B parameters across the days range and growing to3434B and5858B parameters at the two largest node budgets\. The direction of this trend is that larger hardware budgets support larger active parameter counts, though the near\-ties at neighboring budgets mean the exact crossover points are not resolved at this granularity\. These findings suggest that an interior optimum in sparsity appears once the systems constraints are added, where a memory floor bounds total parameters and a GPU\-hour cost charges for the MFU degradation and all\-to\-all penalties of highSSand highGG\. We note that this avoids any explicit cap on model dimensions, since the systems constraints alone bound the feasible region and let an interior sparsity emerge\.

### 7\.5From Predicted Ranking to Measured Loss

In this final experiment, we test whether the configuration ordering predicted by MOSAIC transfers to real training runs\. We take a small budget envelope of eight p6 nodes \(6464GPUs\) for1010days, enumerate the feasible grid, and select four configurations from the ranked predictions\. These four training configurations span the same trade\-off MOSAIC describes, and we launch each as a real pretraining run\. Table[5](https://arxiv.org/html/2608.10605#S7.T5)lists the selected configurations with the MFU the framework predicts for each\.

rank \(loss\)LlayersL\_\{\\text\{layers\}\}EEKKGGNactN\_\{\\text\{act\}\}\(B\)NtotN\_\{\\text\{tot\}\}\(B\)SSpred\. MFU \(%\)norm\. actual MFUppepcptpmbsID 140128248\.5250\.00\.96614\.00\.8448114ID 240128285\.6126\.40\.95611\.60\.6948114ID 34032248\.566\.00\.87221\.80\.9148114ID 43232227\.365\.80\.88923\.31\.048114

Table 5:The four configurations on the88\-node,1010\-day envelope, ranked by predicted loss\. We keep the global batch size fixed to 1024 across all runs and assign each configuration’s peak learning rate using Equation \([17](https://arxiv.org/html/2608.10605#A1.E17)\) in Appendix[A\.2](https://arxiv.org/html/2608.10605#A1.SS2)\. The normalized actual MFU column reports each run’s measured MFU divided by the largest measured MFU among these four runs\.The first test is whether the performance model orders these configurations the way the hardware does\. Our measured MFU reproduces this ordering exactly, where the rank correlation between predicted and observed MFU across the four runs aligns as shown in Table[5](https://arxiv.org/html/2608.10605#S7.T5)\. Although MOSAIC relies on the absolute MFU values, we discuss the MFU ordering here to test the selection step itself, since these four configurations were chosen by a predicted ranking rather than by any single MFU number, and we already established in Section[7\.2\.1](https://arxiv.org/html/2608.10605#S7.SS2.SSS1)the absolute accuracy on which the feasibility constraint and the token budget rest\. Normalizing the predicted MFU the same way gives0\.600\.60,0\.500\.50,0\.940\.94, and1\.001\.00against the measured0\.840\.84,0\.690\.69,0\.910\.91, and1\.001\.00, and we observe that the predicted spread is wider than the measured one even though the ordering agrees\.

We next examine the loss trajectory of the runs themselves, but where we plot against both model FLOPs axis and an axis that accounts for the efficiency in MFU\. The eight\-node, ten\-day envelope is only the selection envelope MOSAIC ranks against, and the staged runs train beyond it\. The hardware\-compute axis of Figure[12](https://arxiv.org/html/2608.10605#S7.F12)divides by the measured MFU of each run\.

The largest model FLOPs these curves reach is roughly1515–20×20\\timesthe largest budget in the scaling law fitting set of Appendix[C](https://arxiv.org/html/2608.10605#A3), which tops out at8\.1×10208\.1\\times 10^\{20\}FLOPs, so the runs probe the predicted ordering well outside the regime the law was fit on\. We plot the loss of the four runs against model FLOPs and against the peak\-equivalent hardware compute consumed in Figure[12](https://arxiv.org/html/2608.10605#S7.F12), with the warm\-up discarded\. ID 3 and ID 4 isolate the expert split factor, since they sit at nearly the same total parameters \(66\.066\.0against65\.865\.8B\) and close sparsity \(0\.8720\.872against0\.8890\.889\), yet theG=4G\\\!=\\\!4run stays below theG=2G\\\!=\\\!2run on both axes\. On the per\-FLOP axis, up to roughly102210^\{22\}model FLOPs theG=8G\\\!=\\\!8, low\-NactN\_\{\\text\{act\}\}configuration \(ID 2\) holds an advantage over ID 1, since its smaller active count \(5\.65\.6B against77–8\.58\.5B\) shifts its curve toward fewer model FLOPs to reach a given loss\. Past that point the two curves converge and the margin narrows to the width of the smoothed trajectories, so we read the advantage as a tendency over the budget range we ran rather than a separation that holds at the end of training\. Even so, the direction matches the compute\-frontier analysis, where the run slowest in wall\-clock terms converts a fixed FLOPs budget into loss most efficiently\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/x7.png)Figure 12:Smoothed loss trajectories of the staged runs against model FLOPsCmodelC\_\{\\text\{model\}\}\(left\) and against the peak\-equivalent hardware compute consumedCmodel/MFUC\_\{\\text\{model\}\}/\\mathrm\{MFU\}\(right\), with the warm\-up discarded\. ID 2 leads over much of the lower model\-FLOPs range, while the trajectories converge and ID 1 slightly overtakes it near the end \(at102210^\{22\}FLOPs\) within the width of the smoothed curves\. On the hardware\-compute axis, the advantage of ID 1 over ID 2 appears much earlier\.The observation in Figure[12](https://arxiv.org/html/2608.10605#S7.F12)changes once the same trajectories are placed on a hardware\-compute axis rather than a model\-FLOPs one\. We use the peak\-equivalent hardware compute consumed,Ccost=Cmodel/\(MFU​ηgood\)C\_\{\\text\{cost\}\}=C\_\{\\text\{model\}\}/\(\\mathrm\{MFU\}\\,\\eta\_\{\\text\{good\}\}\), the compute proportional to the GPU\-hours a run spends to deliverCmodelC\_\{\\text\{model\}\}, and we holdηgood=1\\eta\_\{\\text\{good\}\}=1throughout, soCcostC\_\{\\text\{cost\}\}reduces toCmodel/MFUC\_\{\\text\{model\}\}/\\mathrm\{MFU\}as plotted\. Dividing out MFU charges each run for the hardware time it actually spends, and ID 1 sustains the higher MFU \(normalized0\.840\.84against ID 2’s0\.690\.69\), so the same hardware buys ID 1 about1\.2×1\.2\\timesthe model computation\. Figure[12](https://arxiv.org/html/2608.10605#S7.F12)shows this directly, where ID 2 leads on the model\-FLOPs panel \(left\) while ID 1 leads on theCcostC\_\{\\text\{cost\}\}panel \(right\) toward the longer horizon of the run\.

Taken together, the staged runs provide two observations\. First, the predicted MFU ranking reproduces on hardware, so the selection that MOSAIC relies on transfers from those predictions\. Second, the ordering between the two leading configurations changes between the model\-FLOPs and hardware\-cost views, which is the co\-design effect MOSAIC is designed to capture\. We report these as an early validation of our framework\. The results above optimize one hardware envelope at a time\. In Appendix[D\.3](https://arxiv.org/html/2608.10605#A4.SS3), we ask whether the discrete MOSAIC optima can be summarized by a Chinchilla\-style optimal resource\-allocation law that maps a purchased hardware budget to the jointly optimal model and data allocation\. For a fixed architecture, the reduction is exact\. The hardware budget determines its optimal token count and therefore its predicted loss\. This also gives a simple phase structure, where the preferred architecture can change as the training budget grows\. So we present the analysis as a direction for future work rather than as a main empirical result\.

## 8Related Work

##### MoE Scaling Laws\.

Past work on sparse MoE architecture extends Chinchilla\-style scaling laws\(hoffmann2022training\)to sparse MoEs\.krajewski2024scalingintroduce granularity as an explicit axis andabnar2025parameterscharacterize the sparsity dimension, whileludziejewski2025jointfit a joint law over active size, experts, and tokens and argue MoEs can be more memory\-efficient than dense models\.wan2026holisticpush toward holistic recipes that map a budget to a full architecture, andmcleish2025gemstonesshow how architecture and hyperparameter choices change the prescriptions a scaling law fit yields\. Our scaling law work extends these previous studies but isolates the MoE\-layer geometry and is coupled to a systems cost model rather than read under FLOPs alone\. There have been few more notable studies in the arena of scaling laws for MoE models that focus on multiple dimensions\(tian2026towards\), however, as discussed in our desiderata, we focus on low dimensional MoE scaling laws in this work that suit the systems derivation\.

##### Systems Performance Estimation\.

Measuring realized throughput of distributed training rather than its FLOPs is central to the thesis of MOSAIC\. While this usually falls under systems performance engineering in most works, we focus mainly on modeling and not on engineering the systems like writing better kernels\. To that end,fernandez2024hardwaremeasure diminishing returns and parallelism strategy crossovers empirically, while analytical co\-design tools such as Calculon\(isaev2023calculon\)and profiling driven simulators such as vTrain\(bang2024vtrain\)predict training time and memory from a model and layout description\.bian2025inferencetarget inference\-efficient model selection, and on the MoE sidehuang2023towardsandzheng2026uniepattack expert dispatch, all\-to\-all, and kernel inefficiencies in deployment and training\. These efforts model or optimize execution cost in isolation\. We instead fold an operator\-level MFU and memory predictor directly into the architecture\-selection objective\.

##### Joint Architecture\-Systems Co\-Design\.

Finally, we discuss the relevant studies that optimize model and hardware together\.anthony2024codesignalign model dimensions to GEMM and wave\-quantization boundaries whereasbian2026architecturejointly tune architecture for accuracy and inference efficiency\.bae2025hybridstudies hybrid architectures for compute optimality that are better on recall andwang2025catransformersco\-design transformers and accelerators under a carbon objective\. We also see studies likezhang2025ladderwhich redefine tensor parallelism via a Ladder Residual architectural change that overlaps communication with computation for accelerated inference, a complementary form of architecture\-systems co\-design that adapts the model to the parallel\-execution substrate rather than the GEMM tile\. We also find that classic budget studies such aslescao2022millionablate architecture at fixed GPU\-hours, andwu2024inferenceadd an inference\-compute axis\. We differ by coupling a predictive MoE scaling law to a parallelism\-aware throughput model, recovering deployable MoE geometries under explicit memory and GPU\-hour constraints rather than FLOPs\.

## 9Conclusion, Limitations and Future Work

We introduced MOSAIC, a systems\-aware scaling framework that jointly selects the model architecture, training\-token budget, and distributed parallel layout for a fixed cluster and training window\. Using roughly 150 from\-scratch sparse MoE pretraining runs, we fit a joint scaling law over total parameters, sparsity, training tokens, and expert split factor, and show that model\-FLOPs only optimization admits no interior sparsity optimum within the calibrated sparsity region\. To enable the joint systems optimization, we developed an operator\-level performance model that predicts MFU, memory, and the best parallel layout across models ranging from 700 million to 18 billion active parameters, while reliably recovering the highest\-throughput configurations\. We observe that coupling the scaling law and performance model changes the architecture prescription\. Notably, the optimal sparsity becomes interior and hardware\-dependent rather than following the monotone boundary preference implied by model FLOPs, while the expert geometry shifts with the cluster size and training horizon\. Our budget sweeps recover interior, hardware\-dependent MoE configurations, and our staged pretraining runs reproduce the predicted MFU ordering on real hardware, with the loss ordering of the runs flipping between the model\-FLOPs axis and the hardware\-compute axis\.

One limitation of MOSAIC is that the precise architectures returned are conditional on the chosen geometry ladder, hardware platform, and training recipe, and some reported optima require substantial extrapolation beyond the scaling law training regime\. Moreover, uncertainty in the non\-identified scaling law coefficients is not propagated through the discrete architecture search\. Consequently, our evidence more strongly supports the hardware\-aware ranking and co\-design principle than any single geometry as a universal cluster optimum\. In addition, our performance model mirrors a fairly basic Megatron\-Core training stack, and the absence of a clear monotone trend in the optimal expert split factorGGmay itself be an artifact of that stack’s systems penalty\. Thus, broader kernel\-level changes could turn the optimalGGinto a predictable function of the budget, which we consider both a limitation and grounds for future expansive validations\. Finally, our runs limit the expert split factor toG≤8G\\leq 8, and future work could study extrapolations of the law along that dimension\. Our studies are also carefully designed to bound the tokens\-per\-parameter ratio, so we do not study the overtrained regime where models train far beyond these ratios, and we leave that as future work\. For future work, we layout the following directions that we believe are worth pursuing\.

- •Joint batch\-size and parallelism optimization\.In this work, we fixed the global batch size across all our runs\. Extending the decision space to include it would allow MOSAIC to jointly reason about convergence, strong scaling, and the layout changes induced by the number of microbatches in flight\.
- •Broader and accelerator\-aligned geometry families\.Our current scaling law sweep and MOSAIC search share a single backbone geometry ladder\. Extending the fit across multiple depth, width anddf​fd\_\{ff\}co\-scaling families would test whether the\(Ntot,S,G\)\(N\_\{\\mathrm\{tot\}\},S,G\)representation transfers across backbone geometries and would allow MOSAIC to search beyond the family used to fit the law\. Incorporating hardware\-friendly tile and wave constraints into these families would further enable accelerator\-specific architecture selection\.
- •Communication and computation overlap\.Our current performance model costs compute and exposed communication largely additively, overlapping only the data\-parallel collectives where scheduling windows allow\. Extending the performance model to represent persistent and fused kernels that overlap collectives with computation would improve MFU prediction as modern training stacks increasingly depart from additive compute\-and\-communication execution\.

Acknowledgments and Disclosure of Funding

We thank Valentyn Boreiko, Moritz Haas, Srikanth Doss, Volkan Cevher, Cole Hawkins, Leonard Lausen, Anna Ruminshky, Hongshan Li and Subendhu Rongali for their help in setting up the infrastructure and the training pipeline used to train the models in this paper\. We also thank Justin Chiu for the initial version of the analytical model implementation used in this paper\.

## References

## Appendix APre\-training Recipe

In this section, we cover some more details of the pretraining recipe and the architecture that went into the training runs\.

### A\.1Architecture and Pretraining Details

##### Sparse MoE Architecture\.

Our MoE block follows the DeepSeekMoE design\(dai2024deepseekmoe\), which splits the reference FFN into fine\-grained routed experts and adds a shared expert that every token passes through, paired with the auxiliary\-loss\-free bias balancing ofwang2024lossfreein the router\. In our instantiation, the shared expert has the same hidden widthdexpertd\_\{\\text\{expert\}\}as a routed expert, so as the expert split factorGGincreases,dexpertd\_\{\\text\{expert\}\}decreases and so does the shared expert size\. The attention mechanism is gated\(qiu2025gated\)and a head\-specific sigmoid gate is applied to the output of each scaled\-dot\-product attention head before the output projection, which the authors find consistently improves loss for both dense and MoE models at this scale\. We use pre\-LayerNorm and SwiGLU activations inside each expert and the embedding and unembedding matrices are not tied\. We do not use QK layernorm for the attention layer\.

##### Pretraining objective and data\.

Every run is a pure text\-to\-text pretraining run with the next\-token\-prediction objective of GPT\-2\(radford2019language\)\. We do not use any auxiliary denoising or instruction\-tuning loss aside the sequence auxiliary loss in the MoE routing layer which we discuss in later sections\. The training corpus is the English portion of mC4\(xue2020mt5\), the multilingual extension of the Colossal Clean Crawled Corpus \(C4\) that is itself a cleaned subset of Common Crawl\. We tokenize with the GPT\-2 BPE tokenizer, which has a vocabulary of 50,257 tokens\.

##### Numerical precision\.

All training runs use bf16 weights and activations with a fp32 master copy of the optimizer state\.

##### Initialization\.

We use the Megatron/GPT\-2 architecture in our runs\(radford2019language\)\. Fan\-in \(input\-side\) projections, the QKV and first FFN matrices, the routerWrW\_\{r\}, and the embedding and readout matrices, are drawn flat from𝒩​\(0,0\.0062\)\\mathcal\{N\}\(0,0\.006^\{2\}\)\. Output\-side residual projections, the attention output projection and the expert/FFN down\-projection, are depth\-scaled to𝒩​\(0,\(0\.006/2​Llayers\)2\)\\mathcal\{N\}\\\!\\bigl\(0,\(0\.006/\\sqrt\{2L\_\{\\text\{layers\}\}\}\)^\{2\}\\bigr\), so deeper models initialize the residual\-write weights smaller and keep the residual\-stream variance stable across the2​Llayers2L\_\{\\text\{layers\}\}residual additions per forward pass\. For the normalization layers, their gains are set to11and biases to0\. InitializingWrW\_\{r\}like any other fan\-in matrix leaves router logits approximately uniform across experts at step zero\. We leaveμ\\muP\-style initialization for hyperparameter transfer\(jiang2026hpmoe\)in sparse MoE models as future work\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/loss_vs_lr.png)\(a\)Loss vs\. peak learning rate per configuration, with the per\-configuration optimum \(stars\) read from a quadratic fit inlog10⁡lr\\log\_\{10\}\\mathrm\{lr\}\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/lr_vs_compute_band.png)\(b\)Optimal learning rate vs\. model FLOPsCmodel=6​Nact​DC\_\{\\text\{model\}\}=6\\,N\_\{\\text\{act\}\}\\,D, solid line is the fitlr⋆​\(Cmodel\)=3\.30​Cmodel−0\.199\\mathrm\{lr\}^\{\\star\}\(C\_\{\\text\{model\}\}\)=3\.30\\,C\_\{\\text\{model\}\}^\{\-0\.199\}, shaded region is the near\-optimal band\.

Figure 13:Setting the peak learning rate\. Left: loss\-vs\-learning\-rate curves across MoE configurations, each with a quadratic\-fit optimum\. Right: those optima collapse onto a single power law in the model FLOPsCmodelC\_\{\\text\{model\}\}, with one law predicting the optimal learning rate across the full range of total parameters and sparsity\.

### A\.2Optimizer and Learning\-rate Settings

We use AdamW at weight decay0\.10\.1with the warmup\-stable\-decay \(WSD\) schedule ofhagele2024scaling, namely a fixed warmup of approximately100100\-million tokens regardless of run length, a stable phase at peak learning rate, and a linear decay to zero over the final20%20\\%training steps\. The AdamW moments are the standard defaultsβ1=0\.9\\beta\_\{1\}=0\.9andβ2=0\.999\\beta\_\{2\}=0\.999withε=10−14\\varepsilon=10^\{\-14\}, and we do not clip gradients\. We do not conduct an independent learning\-rate sweep for every scaling law run\. Instead, each run’s peak learning rate is assigned from a single compute\-conditioned rule fitted on a separate pilot sweep\. To keep learning rate from confounding the scaling law runs, we set the peak learning rate from a single power law in the model FLOPsCmodel=6​Nact​DC\_\{\\text\{model\}\}=6\\,N\_\{\\text\{act\}\}\\,D\. We sweep a small grid of peak learning rates across MoE configurations spanning a wide range of total parameters, sparsity, and token budgets\. Each configuration traces a loss\-vs\-learning\-rate curve whose optimum we read from a quadratic fit inlog10⁡lr\\log\_\{10\}\\mathrm\{lr\}\(Figure[13](https://arxiv.org/html/2608.10605#A1.F13), left\)\. For each configuration we keep the near\-optimal band, every swept rate whose loss lands within a small tolerance of that configuration’s best, following the band\-fit practice of prior learning\-rate studies\(tian2026towards\)\. Fittinglog10⁡lr⋆\\log\_\{10\}\\mathrm\{lr\}^\{\\star\}againstlog10⁡Cmodel\\log\_\{10\}C\_\{\\text\{model\}\}over these near\-optimal points gives

lr⋆​\(Cmodel\)=3\.30​Cmodel−0\.199,Cmodel=6​Nact​D\\mathrm\{lr\}^\{\\star\}\(C\_\{\\text\{model\}\}\)\\;=\\;3\.30\\;C\_\{\\text\{model\}\}^\{\-0\.199\},\\qquad C\_\{\\text\{model\}\}=6\\,N\_\{\\text\{act\}\}\\,D\\,\(17\)with a tight parallel band around the central line \(Figure[13\(b\)](https://arxiv.org/html/2608.10605#A1.F13.sf2)\) where any rate inside the band is well tuned, and the band width reflects the resolution of the swept grid\.

We use compute alone as the axis for learning rate predictions and we leave more appropriate sparsity specific learning rate predictions \(if there are any\) for future work\. AWithin our pilot sweep, after conditioning onCmodelC\_\{\\mathrm\{model\}\}, we detect no systematic residual dependence of the near\-optimal learning rate onNtotN\_\{\\mathrm\{tot\}\}, sparsity, or theNact/DN\_\{\\mathrm\{act\}\}/Dsplit\. We therefore use the compute\-conditioned rule rather than fitting a higher\-dimensional learning\-rate surface, which would require substantially more dedicated hyperparameter sweeps and could itself confound the scaling law fit\. The decay exponent also agrees with independent MoE learning\-rate fits, wheretian2026towardsreport a slope of−0\.153\-0\.153over a comparable compute range, close to our value\. In practice, this gives a single rule, where for any candidate configuration we computeCmodel=6​Nact​DC\_\{\\text\{model\}\}=6\\,N\_\{\\text\{act\}\}\\,Dand read the peak learning rate off Equation \([17](https://arxiv.org/html/2608.10605#A1.E17)\)\.

##### Out\-of\-scope HPO axes\.

Optimizer hyper\-parameters \(weight decay0\.10\.1, warmup, scheduler shape\), sequence length \(Tseq=2048T\_\{\\text\{seq\}\}=2048\), drop\-token capacity factor \(default infinite capacity\), are each held to a single value across all runs\. The law is therefore conditional on this recipe\.

### A\.3MoE Routing Details

In this section, we describe the choices behind the routing algorithm for our scaling recipe presented in Algorithm[1](https://arxiv.org/html/2608.10605#alg1)\.

##### Sigmoid scoring with per\-expert bias\.

Followingwang2024lossfree, we instantiate the score vector𝒔​\(x\)\\bm\{s\}\(x\)of Equation \([7](https://arxiv.org/html/2608.10605#S3.E7)\) as a clean sigmoid scoresi′​\(x\)s^\{\\prime\}\_\{i\}\(x\)plus an additive per\-expert biasϕ∈ℝE\\phi\\in\\mathbb\{R\}^\{E\}, so that

si​\(x\)=si′​\(x\)\+ϕi=σ​\(θi​\(x\)\)\+ϕi,θi​\(x\)=w~r,i⊤​xs\_\{i\}\(x\)\\;=\\;s^\{\\prime\}\_\{i\}\(x\)\+\\phi\_\{i\}\\;=\\;\\sigma\\\!\\bigl\(\\theta\_\{i\}\(x\)\\bigr\)\+\\phi\_\{i\},\\qquad\\theta\_\{i\}\(x\)\\;=\\;\\tilde\{w\}\_\{r,i\}^\{\\\!\\top\}\\,x\(18\)wherew~r,i\\tilde\{w\}\_\{r,i\}denotes the row\-normalizedii\-th row ofWrW\_\{r\}\. The sigmoid form decouples per\-expert logits, avoiding the shared\-softmax competition that drives unbounded logit growth and forces ad\-hoc remedies such as routerzz\-loss or logit capping\. The biasϕ\\phiis not a learned parameter and it is updated by a controller using only routing statistics, and it influences which experts are selected\.

##### LFLB bias controller\.

Following Algorithm[1](https://arxiv.org/html/2608.10605#alg1), letnin\_\{i\}denote the number of tokens routed to expertiiin a batch and letn¯=1E​∑i=1Eni\\bar\{n\}=\\frac\{1\}\{E\}\\sum\_\{i=1\}^\{E\}n\_\{i\}\. Followingwang2024lossfree;team2025kimi, we updateϕ\\phiin the direction of imbalance, but adopt an RMS\-normalized violation rather than a raw sign so that the magnitude of the correction is proportional to the shape of the imbalance pattern and not the absolute batch size, as follows

ui=n¯−niE−1​∑j\(n¯−nj\)2\+ε,ϕi←ϕi\+μb​uiu\_\{i\}\\;=\\;\\frac\{\\bar\{n\}\-n\_\{i\}\}\{\\sqrt\{E^\{\-1\}\\sum\_\{j\}\(\\bar\{n\}\-n\_\{j\}\)^\{2\}\+\\varepsilon\}\},\\qquad\\phi\_\{i\}\\;\\leftarrow\\;\\phi\_\{i\}\+\\mu\_\{b\}\\,u\_\{i\}\(19\)We additionally center the bias after each update,ϕi←ϕi−1E​∑jϕj\\phi\_\{i\}\\leftarrow\\phi\_\{i\}\-\\frac\{1\}\{E\}\\sum\_\{j\}\\phi\_\{j\}, since only relative bias differences affect routing\.

##### Probabilistic top\-KKvia Gumbel proposal–correction\.

We replace deterministic selection with a Gumbel\-noise proposal–correction sampler operating at a fixed temperatureτ\\tau:

\(Gumbel proposal\)𝒞t\\displaystyle\\text\{\(Gumbel proposal\)\}\\quad\\mathcal\{C\}\_\{t\}=Top⁡Mi​\(log⁡st,i′\+τ​gt,i\),\\displaystyle=\\operatorname\{Top\}\\\!M\_\{i\}\\bigl\(\\log s^\{\\prime\}\_\{t,i\}\+\\tau\\,g\_\{t,i\}\\bigr\),gt,i\\displaystyle g\_\{t,i\}∼Gumbel​\(0,1\),\\displaystyle\\sim\\mathrm\{Gumbel\}\(0,1\),\(20\)\(Gumbel correction\)𝒯K,t\\displaystyle\\text\{\(Gumbel correction\)\}\\quad\\mathcal\{T\}\_\{K,t\}=Top⁡Ki∈𝒞t​\(st,i\),\\displaystyle=\\operatorname\{Top\}\\\!K\_\{i\\in\\mathcal\{C\}\_\{t\}\}\\bigl\(s\_\{t,i\}\\bigr\),M\\displaystyle M=min⁡\(ρ​K,E\)\.\\displaystyle=\\min\(\\rho K,E\)\.\(21\)The construction rests on the Gumbel\-max theorem\. ForK=1K\\\!=\\\!1andgi∼Gumbel​\(0,1\)g\_\{i\}\\sim\\mathrm\{Gumbel\}\(0,1\),

arg​maxi⁡\(log⁡pi\+τ​gi\)∼Categorical​\(pi1/τ∑jpj1/τ\)\\operatorname\*\{arg\\,max\}\_\{i\}\\bigl\(\\log p\_\{i\}\+\\tau\\,g\_\{i\}\\bigr\)\\;\\sim\\;\\mathrm\{Categorical\}\\\!\\left\(\\tfrac\{p\_\{i\}^\{1/\\tau\}\}\{\\sum\_\{j\}p\_\{j\}^\{1/\\tau\}\}\\right\)\(22\)and its top\-KKextension gives exact sampling without replacement from the same distribution\(yellott1977relationship;kool2019stochastic\)\. Injecting the noise in log\-score space, as in Equation \([20](https://arxiv.org/html/2608.10605#A1.E20)\), therefore draws the proposal set exactly fromP​\(i\)∝st,i′⁣1/τP\(i\)\\propto s^\{\\prime 1/\\tau\}\_\{t,i\}on the clean scores, and the biasϕ\\phiis then applied in the correction stage of Equation \([21](https://arxiv.org/html/2608.10605#A1.E21)\) to preserve those sampling semantics\. This tempered sampling interpretation holds for the proposal before centering and correction, and the implemented algorithm modifies that distribution\. The per\-expert centering of Algorithm[1](https://arxiv.org/html/2608.10605#alg1)makes thegt,ig\_\{t,i\}dependent across the batch and the bias\-aware correction reweights the survivors, so we treat both as heuristic variance\-control and load\-balancing steps rather than as sampling from the tempered clean\-score distribution\. The proposal size isM=min⁡\(ρ​K,E\)M=\\min\(\\rho K,E\)withK≤EK\\leq E, sinceρ​K\\rho KexceedsEEat the small expert counts in our grid\. The temperatureτ\\tausets how peaked the proposal distribution is, and asτ→0\\tau\\\!\\to\\\!0the algorithm approaches a deterministic clean\-score top\-MMproposal followed by the bias\-aware top\-KKcorrection within that proposal\. This limit need not coincide with the global deterministic top\-KKonst,is\_\{t,i\}, since the bias can favor an expert outside the clean top\-MM\. For the uncentered, uncorrected Gumbel\-max baseline with two experts andK=1K\\\!=\\\!1, the probability of selectingbboveraaisσ​\(τ−1​log⁡\(sb′/sa′\)\)\\sigma\\\!\\bigl\(\\tau^\{\-1\}\\log\(s^\{\\prime\}\_\{b\}/s^\{\\prime\}\_\{a\}\)\\bigr\), so near\-tied experts receive meaningful exploration while confident decisions are exponentially protected\. In Algorithm[1](https://arxiv.org/html/2608.10605#alg1)itself, at two expertsρ=12\\rho=12puts both in the proposal and the correction then selects deterministically underst,is\_\{t,i\}\.

As a final, empirically important step, we center the noise per expert across the current batch,g~t,i=gt,i−1T​∑tgt,i\\tilde\{g\}\_\{t,i\}=g\_\{t,i\}\-\\frac\{1\}\{T\}\\sum\_\{t\}g\_\{t,i\}, so that finite\-sample noise means do not waste the bias controller’s correction budget chasing purely stochastic imbalances\. We setτ=0\.01\\tau=0\.01andρ=12\\rho=12, with token\-level granularity across all our runs, and we useμb=10−3\\mu\_\{b\}=10^\{\-3\}andε=10−4\\varepsilon=10^\{\-4\}in the bias update of Equation \([19](https://arxiv.org/html/2608.10605#A1.E19)\)\. We find that this method of sampling with lowτ\\tauallows us to overcome issues of load balancing and experts starved of tokens in a single batch later in long horizon runs especially at large active parameters of the scale of 18 billion parametersNactN\_\{\\text\{act\}\}and more\. Compared to adding more dense layers at the beginning as done insingh2026arceewhich reduces the model capacity at matchedNactN\_\{\\text\{act\}\}, we find that our Gumbel sampling method along with the proposed routing algorithm helps in more efficient mitigation of instabilities due to spikes in starved experts\.

##### Router\-weight row normalization\.

A subtle source of routing pathology is heterogeneity in the row norms ofWrW\_\{r\}\. Experts whose corresponding rows have grown a larger norm are favored for reasons unrelated to specialization\. We optionally row\-normalize the router weights at every forward passw~r,i=cr⋅wr,i‖wr,i‖2\\tilde\{w\}\_\{r,i\}\\;=\\;c\_\{r\}\\cdot\\frac\{w\_\{r,i\}\}\{\\\|w\_\{r,i\}\\\|\_\{2\}\}, where the target normcrc\_\{r\}is a fixed constant, in our case set to 1\.0\. It allows to recover routing decisions that depend on directional alignment betweenwr,iw\_\{r,i\}andxxrather than on accidental norm drift\. This is theRowNorm​\(Wr,cr\)\\textsc\{RowNorm\}\(W\_\{r\},c\_\{r\}\)step in Algorithm[1](https://arxiv.org/html/2608.10605#alg1), which produces the normalized routerW~r\\tilde\{W\}\_\{r\}used to compute the pre\-sigmoid logitsΘ\\Theta\.

Additionally, the bias update of Eq\. \([19](https://arxiv.org/html/2608.10605#A1.E19)\) is invoked once per optimizer step on the per\-expert countsnnreturned by the algorithm\. We emphasize that load balancing is carried by the bias controller and the noise distribution, where no batch\-level load\-balancing loss term is added to the optimization objective, in contrast to the auxiliary LBL ofshazeer2017outrageously;fedus2022switch\. The one auxiliary term we retain is the sequence\-level loss described below\.

##### Router regularization\.

In addition to the LFLB bias controller of Section[3\.3](https://arxiv.org/html/2608.10605#S3.SS3), we apply the routerzz\-loss ofzoph2022stto suppress unbounded growth in the pre\-sigmoid router logitsΘt,i=θi​\(xt\)\\Theta\_\{t,i\}=\\theta\_\{i\}\(x\_\{t\}\)of Eq\. \([18](https://arxiv.org/html/2608.10605#A1.E18)\):

ℒz=1T​∑t=1T\(log​∑i=1EeΘt,i\)2\\mathcal\{L\}\_\{z\}\\;=\\;\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\Bigl\(\\log\\sum\_\{i=1\}^\{E\}e^\{\\Theta\_\{t,i\}\}\\Bigr\)^\{2\}\(23\)added to the cross\-entropy loss with a small fixed coefficient\. Because selection in our routing variant is taken onσ​\(Θ\)\+ϕ\\sigma\(\\Theta\)\+\\phirather than onΘ\\Thetadirectly,ℒz\\mathcal\{L\}\_\{z\}acts purely as a numerical guard on the logits and does not interfere with the bias controller\. We set the coefficient ofzz\-loss to 1e\-4 in all our runs\.

##### Sequence Auxiliary Loss\.

Unlike the load balancing batch auxiliary loss used in the softmax routing variants, we follow the sequence auxiliary loss as used indeepseekai2024v3\. We set its coefficient to 1e\-3 in all our runs\.

## Appendix BScaling Analysis

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/G_geometry.png)Figure 14:Example of an MoE layer in a transformer block, at three expert split factorsG∈\{1,4,8\}G\\\!\\in\\\!\\\{1,4,8\\\}\. In every panel, the routed experts sit beside a shared expert\.dffd\_\{\\text\{ff\}\}is the hidden width of a single reference dense feed\-forward expert \(G=1G\\\!=\\\!1\), and the total routed model widthE​dffE\\,d\_\{\\text\{ff\}\}isheld fixedacross the panels\. RaisingGGsplits each reference FFN intoGGnarrower experts, so the routed\-expert count grows \(E→G​EE\\\!\\to\\\!G\\,E\) while each expert’s hidden width shrinks proportionally \(dexpert=dff/Gd\_\{\\text\{expert\}\}=d\_\{\\text\{ff\}\}/G\)\. Relative toG=1G\\\!=\\\!1, a factor\-GGconfiguration hasGGtimes as many experts, eachGGtimes narrower\. The top\-KKselection is kept the same across the three panels, so the active parametersNactN\_\{\\text\{act\}\}shrink with higherGGrather than staying fixed\. Note that this is different from granularity defined inkrajewski2024scaling\.This appendix collects the per\-GGevidence behind the scaling laws of Section[4](https://arxiv.org/html/2608.10605#S4), starting with Figure[14](https://arxiv.org/html/2608.10605#A2.F14), which shows how the expert split factorGGof Section[3](https://arxiv.org/html/2608.10605#S3)reshapes an MoE layer\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_G1.png)\(a\)G=1G=1\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_G2.png)\(b\)G=2G=2\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_G4_main.png)\(c\)G=4G=4\.
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/parity_G8.png)\(d\)G=8G=8\.

Figure 15:Predicted\-vs\-observed loss for theℒ​\(Ntot,S,D\)\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D\)fit with Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\), refit independently on each expert split factor stratumG∈\{1,2,4,8\}G\\\!\\in\\\!\\\{1,2,4,8\\\}\. TheG=4G\\\!=\\\!4andG=8G\\\!=\\\!8panels are the ones reproduced in Figure[3](https://arxiv.org/html/2608.10605#S4.F3)\.### B\.1Per\-GGScaling Law Fits

This appendix section gives the per\-GGdetails behind Section[4\.1](https://arxiv.org/html/2608.10605#S4.SS1)\. We fitℒ​\(Ntot,S,D\)\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D\)in Equation \([9](https://arxiv.org/html/2608.10605#S4.E9)\) on eachGGseparately, holding out the largest10%10\\%by active parameters\. The fit is noise\-limited in all the perGGoptimization of the functional form and the held\-out large models match it, with no residual structure onEE,KK,NtotN\_\{\\text\{tot\}\},SS,DD, orLlayersL\_\{\\text\{layers\}\}\(all Pearson correlations below0\.130\.13in magnitude\)\. Refitting the same form independently per expert split factor as shown in Figure[15](https://arxiv.org/html/2608.10605#A2.F15), \(G∈\{1,2,4,8\}G\\\!\\in\\\!\\\{1,2,4,8\\\}\) fits eachGGwell in isolation, but the fitted coefficients drift substantially acrossGG, so no single set of\(Ntot,S,D\)\(N\_\{\\text\{tot\}\},S,D\)coefficients describes all expert split factors at once\. This is the signal thatGGmust enter the law itself, which we do in Section[4\.2](https://arxiv.org/html/2608.10605#S4.SS2)\.

### B\.2Joint\-Law Uncertainty and Identifiability

This appendix section gives the evidence behind theGG\-augmented joint law of Section[4\.2](https://arxiv.org/html/2608.10605#S4.SS2), namely that a larger expert split factorGGlowers loss at a fixed parameter and data budget\. When we fit the joint law, this preference shows up as the split\-factor exponentη≈\+0\.95\\eta\\\!\\approx\\\!\+0\.95, and the question is whether that value is a real effect or an artifact of the particular runs we happened to collect\. To quantify uncertainty inη\\eta, we perform500500warm\-started wild\-bootstrap refits using Rademacher residual perturbations and report a95%95\\%BCa interval\. The resulting interval is\[\+0\.91,\+0\.95\]\[\+0\.91,\+0\.95\], which remains well above zero\.

##### Which coefficients are identified\.

The bootstrap above is warm\-started at the fitted point, so it describes only the local basin around it\. A separate scan of400400fresh restarts shows the loss surface has a wide, near\-flat identifiability ridge along which the prefactors and the remaining exponents move at near\-constant objective\. Onlyη\\etaandβ\\betaare separately identified, the former with the interval reported above and the latter with a BCa interval of\[\+0\.206,\+0\.208\]\[\+0\.206,\+0\.208\]\. The other coefficients are best read as coordinates along this ridge, where their tight local intervals in Table[2](https://arxiv.org/html/2608.10605#S4.T2)understate the true cross\-basin spread, so downstream sensitivity studies should perturb them jointly along the ridge rather than one at a time\.

##### Identifiability caveats\.

Because the fit sits close to the run\-to\-run scatter, the loss surface is flat along a ridge in coefficient space\. The capacity prefactor and exponent\(a,α\)\(a,\\alpha\), the joint\-term parameters\(j,γ,δ\)\(j,\\gamma,\\delta\), and the stand\-alone sparsity floor\(c,λ\)\(c,\\lambda\)trade off against one another at near\-constant objective\. This is why we report onlyη\\etaandβ\\betaas identified and treat the rest as ridge coordinates: the law’s loss value and its per\-GGcalibration are robust, but no single split of the prefactors and exponents is privileged, and pinning one for publication requires fixing a prefactor \(for exampleaa\) and refitting\. The stand\-alone sparsity\-floor termc/\(1−S\)λc/\(1\-S\)^\{\\lambda\}is also numerically small at ourS≤0\.981S\\\!\\leq\\\!0\.981because most of the sparsity penalty is carried by the jointj/\(\(1−S\)δ​Ntotγ​Gη\)j/\(\(1\-S\)^\{\\delta\}\\,N\_\{\\text\{tot\}\}^\{\\gamma\}\\,G^\{\\eta\}\)term\. Finally, the asymptotic blow\-upℒ→∞\\mathcal\{L\}\\\!\\to\\\!\\inftyasS→1S\\\!\\to\\\!1is imposed by the form rather than measured\. Pinning\(c,λ,δ\)\(c,\\lambda,\\delta\)empirically would require runs atSSnear0\.9990\.999, which our grid does not contain\.

## Appendix CScaling Sweep Training Runs

QuantityMinMedianMaxTotal parametersNtotN\_\{\\text\{tot\}\}\(B\)0\.225\.3479\.24Active parametersNactN\_\{\\text\{act\}\}\(B\)0\.1040\.3392\.72SparsitySS0\.4990\.9440\.981Expert split factorGG1\-8Expert countEE8\-512Top\-KKactivation1\-8LlayersL\_\{\\text\{layers\}\}8\-24Residual widthdd1024\-3072Training tokensDD\(B\)2\.885\.7656\.38TPP=D/Nact\\mathrm\{TPP\}=D/N\_\{\\text\{act\}\}10\.8\-70Compute \(non\-embedding FLOPs\)1\.9×10181\.9\\times 10^\{18\}\-8\.1×10208\.1\\times 10^\{20\}Table 6:Coverage of the training run dataset\. These statistics cover only the runs used to fit and validate the scaling law\. The staged MOSAIC validation runs of Section[7\.5](https://arxiv.org/html/2608.10605#S7.SS5), trained for over 200B tokens each, are not included\.GGNtotN\_\{\\text\{tot\}\}\(B\)NactN\_\{\\text\{act\}\}\(B\)SSDD\(B\)LlayersL\_\{\\text\{layers\}\}EE,KK10\.54–30\.910\.152–1\.1790\.72–0\.983\.84–28\.19\{8,12,16\}E∈\{8​–​128\}E\\\!\\in\\\!\\\{8\\text\{\\textendash\}128\\\},K∈\{1,2\}K\\\!\\in\\\!\\\{1,2\\\}20\.95–12\.650\.124–0\.3870\.87–0\.973\.84–5\.76\{8,12,16\}E∈\{32​–​128\}E\\\!\\in\\\!\\\{32\\text\{\\textendash\}128\\\},K∈\{1,2\}K\\\!\\in\\\!\\\{1,2\\\}40\.22–79\.240\.110–2\.7230\.50–0\.982\.88–56\.38\{8,12,16,24\}E∈\{12​–​384\}E\\\!\\in\\\!\\\{12\\text\{\\textendash\}384\\\},K∈\{2,4,6,8\}K\\\!\\in\\\!\\\{2,4,6,8\\\}80\.71–53\.00\.104–2\.520\.85–0\.983\.84–50\.00\{8,12,16,24\}E∈\{96​–​512\}E\\\!\\in\\\!\\\{96\\text\{\\textendash\}512\\\},K∈\{2,4,8\}K\\\!\\in\\\!\\\{2,4,8\\\}Table 7:Ranges of total and active parameters, sparsity, training tokens, depth, expert count, and top\-KKspanned by the scaling runs at each expert split factorGG, where B denotes billions\.### C\.1Coverage of the Scaling\-Law Dataset

This appendix documents the slice of the\(Ntot,Nact,S,G,D\)\(N\_\{\\text\{tot\}\},N\_\{\\text\{act\}\},S,G,D\)space that the joint law of Section[4\.2](https://arxiv.org/html/2608.10605#S4.SS2)is fit on\. The dataset comprisesnruns=150n\_\{\\text\{runs\}\}=150from\-scratch sparse MoE pretraining runs\. Our runs spanNtot∈\[0\.22,79\.24\]N\_\{\\text\{tot\}\}\\in\[0\.22,79\.24\]\\,B parameters,Nact∈\[0\.10,2\.72\]N\_\{\\text\{act\}\}\\in\[0\.10,2\.72\]\\,B parameters,S∈\[0\.499,0\.981\]S\\in\[0\.499,0\.981\],G∈\{1,2,4,6,8\}G\\in\\\{1,2,4,6,8\\\},E∈\{8,12,16,24,32,48,64,96,128,192,256,384,512\}E\\in\\\{8,12,16,24,32,48,64,96,128,192,256,384,512\\\},K∈\{1,2,4,6,8\}K\\in\\\{1,2,4,6,8\\\}, training tokensD∈\[2\.88,56\.38\]D\\in\[2\.88,56\.38\]\\,B \(TPP=D/Nact∈\[10\.8,70\]\\mathrm\{TPP\}=D/N\_\{\\text\{act\}\}\\in\[10\.8,70\]\), across four architecture families \(Llayers,d\)∈\{\(8,1024\),\(12,1536\),\(16,2048\),\(24,3072\)\}\(L\_\{\\text\{layers\}\},d\)\\in\\\{\(8,1024\),\(12,1536\),\(16,2048\),\(24,3072\)\\\}\. Total compute for our scaling law fits spans1\.9×10181\.9\\times 10^\{18\}–8\.1×10208\.1\\times 10^\{20\}FLOPs\. We display some of the statistics of our training runs in Table[6](https://arxiv.org/html/2608.10605#A3.T6)\.

LlayersL\_\{\\text\{layers\}\}ddnrunsn\_\{\\text\{runs\}\}DD\(B\)NtotN\_\{\\text\{tot\}\}\(B\)GGvalues81024522\.88–7\.680\.22–7\.14\{1,2,4,6,8\}\\\{1,2,4,6,8\\\}121536565\.76–11\.520\.73–25\.20\{1,2,4,6,8\}\\\{1,2,4,6,8\\\}1620483928\.19–56\.384\.13–45\.66\{1,2,4,8\}\\\{1,2,4,8\\\}243072450\.0040\.16–79\.24\{4,8\}\\\{4,8\\\}Table 8:Per\-family coverage\.DDis chained to\(Llayers,d\)\(L\_\{\\text\{layers\}\},d\)by an outer Chinchilla\-style budgeting decision\.
### C\.2Training Dynamics of the MoE Models

Figure[16](https://arxiv.org/html/2608.10605#A3.F16)shows the training dynamics of the four largest MoE models in our scaling runs\. For each run we plot, over training steps, the loss, the gradient norm, the number of starved \(under\-utilized\) experts, and the range and mean of the per\-expert routing bias\. We call an expert “starved” when it receives fewer than5%5\\%of the tokens in a batch that a uniform assignment over theEEexperts would send it, so the reported count “starved experts” is the number of such experts in an MoE layer averaged over all layers\. Across all four, the loss descends smoothly, the starved\- expert count collapses to near zero early in training and stays there, and the expert\-bias range contracts steadily, indicating that the load\-balancing controller of Section[3\.3](https://arxiv.org/html/2608.10605#S3.SS3)keeps routing stable and balanced throughout training even at the largest scales\.

![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/training_dynamics_rank1.png)
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/training_dynamics_rank2.png)
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/training_dynamics_rank3.png)
![Refer to caption](https://arxiv.org/html/2608.10605v1/figs/training_dynamics_rank4.png)

Figure 16:Training dynamics of the four largest MoE models in our scaling run sweep \(annotated withNtotN\_\{\\text\{tot\}\},NactN\_\{\\text\{act\}\},EE,KK,GG\)\. Each panel stacks, over training steps, the loss, gradient norm, starved\-expert count, and per\-expert bias range and mean\. Loss descends smoothly, starved experts vanish early, and the bias range contracts, showing stable and balanced routing at scale\.

## Appendix DMOSAIC Optimization

### D\.1Sparse\-MoE Execution\-Plan Realization

For the Megatron\-Core stack used in our experiments, the dense layers factor the world size asNdev=PTP​PCP​PPP​PDPN\_\{\\text\{dev\}\}=P\_\{\\text\{TP\}\}\\,P\_\{\\text\{CP\}\}\\,P\_\{\\text\{PP\}\}\\,P\_\{\\text\{DP\}\}, the factorized device mesh of Megatron\-LM\(shoeybi2019megatron\)and PyTorch\(liang2024torchtitan\), with the degrees running over tensor, context, pipeline, and data parallel groups\. Within an MoE layer, the dense data\-parallel domain is repartitioned into expert parallelism and expert data parallelism,

PDP=PEP​PEDP,PEDP=PDPPEPP\_\{\\text\{DP\}\}\\;=\\;P\_\{\\text\{EP\}\}\\,P\_\{\\text\{EDP\}\},\\qquad P\_\{\\text\{EDP\}\}\\;=\\;\\frac\{P\_\{\\text\{DP\}\}\}\{P\_\{\\text\{EP\}\}\}\(24\)which requires thatPEPP\_\{\\text\{EP\}\}dividesPDPP\_\{\\text\{DP\}\}\. Equivalently, the MoE\-layer mesh factors asNdev=PTP​PCP​PPP​PEP​PEDPN\_\{\\text\{dev\}\}=P\_\{\\text\{TP\}\}\\,P\_\{\\text\{CP\}\}\\,P\_\{\\text\{PP\}\}\\,P\_\{\\text\{EP\}\}\\,P\_\{\\text\{EDP\}\}\. ThusPEPP\_\{\\text\{EP\}\}is an input degree of the plan that participates in the MoE\-layer device factorization, but it does not multiply the dense\-layerPDPP\_\{\\text\{DP\}\}axis independently, and the expert data parallelismPEDPP\_\{\\text\{EDP\}\}is determined rather than chosen\. These constraints, together with the discrete micro\-batch and activation\-checkpointing choices, define the realizable plan set𝒫disc​\(Ndev\)\\mathcal\{P\}\_\{\\text\{disc\}\}\(N\_\{\\text\{dev\}\}\)of Section[2\.2](https://arxiv.org/html/2608.10605#S2.SS2)\.

### D\.2Structural Properties

Three properties of Equation \([5](https://arxiv.org/html/2608.10605#S2.E5)\) are directly useful for interpreting the result\.

##### Property 1 \(objective convexity\)\.

In the continuous variables\(log⁡Ntot,log⁡D,log⁡\(1−S\),log⁡G\)∈ℝ4\(\\log N\_\{\\text\{tot\}\},\\log D,\\log\(1\-S\),\\log G\)\\in\\mathbb\{R\}^\{4\}, the lossℒ​\(Ntot,S,D,G\)\\mathcal\{L\}\(N\_\{\\text\{tot\}\},S,D,G\)of Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) is convex, sinceℒ\\mathcal\{L\}is a posynomial in\(Ntot,D,1−S,G\)\(N\_\{\\text\{tot\}\},D,1\{\-\}S,G\)\(boyd2007tutorial\)on the positive orthant and posynomials are log\-convex in these coordinates\.

##### Property 2 \(constraint\-driven non\-convexity\)\.

The non\-convexity of Equation \([5](https://arxiv.org/html/2608.10605#S2.E5)\) resides in the constraints, not in the objective\.

##### Property 3 \(decoupling ofDDfrom layout\)\.

TokensDDenter only the objective and the FLOPs budget, never the memory cap or the layout product\. Consequently, for anyZZand anyPℓ∈𝒫feas​\(Z,B\)P\_\{\\ell\}\\in\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\), the largest admissibleDDis the closed\-form

D⋆​\(Z,Pℓ\)=Fpeak​Ndev​Ttrain​MFU​\(Z,Pℓ,B\)​ηgood​\(Ndev\)6​Nact​\(Ntot,S\)D^\{\\star\}\(Z,P\_\{\\ell\}\)\\;=\\;\\frac\{F\_\{\\text\{peak\}\}\\,N\_\{\\text\{dev\}\}\\,T\_\{\\text\{train\}\}\\,\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)\\,\\eta\_\{\\text\{good\}\}\(N\_\{\\text\{dev\}\}\)\}\{6\\,N\_\{\\text\{act\}\}\(N\_\{\\text\{tot\}\},S\)\}\(25\)and sinceℒ\\mathcal\{L\}is monotone\-decreasing inDD, any optimum saturates Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\) atD=D⋆​\(Z,Pℓ\)D=D^\{\\star\}\(Z,P\_\{\\ell\}\)\. This eliminates one continuous decision variable analytically and reduces the inner problem to a search over\(Z,Pℓ\)\(Z,P\_\{\\ell\}\)alone\.

##### Algorithmic consequence\.

Properties 1–3 are what let the bilevel program of Equation \([5](https://arxiv.org/html/2608.10605#S2.E5)\) be solved exactly by the single\-stage search of Section[7](https://arxiv.org/html/2608.10605#S7), with no generic MINLP machinery\. The decision space is finitely realizable, since a deployable architecture ranges over the discrete ladder geometries\(Llayers,d,E,K,G\)\(L\_\{\\text\{layers\}\},d,E,K,G\)and an execution plan over small grids of parallelism degrees, micro\-batch sizes, and checkpointing modes, and a fractional expert count or parallel degree is not something one can launch\. Tokens then drop out analytically\. Loss decreases monotonically inDD, so an optimal run never leaves wall\-clock budget unused and Equation \([25](https://arxiv.org/html/2608.10605#A4.E25)\) fixes the saturatingD⋆D^\{\\star\}in closed form, removing the only continuous decision variable\. The deployed search of Section[7](https://arxiv.org/html/2608.10605#S7)realizes this on the finite gridTPP∈\[40,150\]\\mathrm\{TPP\}\\in\[40,150\], taking for each candidate the largest ratio the budget supports, soD⋆D^\{\\star\}is reached up to grid resolution and the cap binds only for the smallest models\. What remains splits cleanly, because at fixedZZthe best execution plan is simply the feasible plan of maximum MFU within Equation \([4](https://arxiv.org/html/2608.10605#S2.E4)\), so the scaling law never participates in the layout search\.

MOSAIC therefore reduces to enumerating the ladder geometries, selecting the max\-MFU feasible plan for each, fixingDDfrom Equation \([25](https://arxiv.org/html/2608.10605#A4.E25)\), and scoring the loss with Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\), and the configuration it returns is the exact optimum of Equation \([5](https://arxiv.org/html/2608.10605#S2.E5)\) over the realizable grid rather than a heuristic solution\. A generic MINLP solver would in any case have little to work with here, sinceMFU\\mathrm\{MFU\}andMem\\mathrm\{Mem\}come out of the calibrated performance model of Section[5](https://arxiv.org/html/2608.10605#S5)as measured kernel and collective lookup tables rather than symbolic expressions that relaxation or branch\-and\-bound machinery can bound\. Searching the discrete geometries directly, rather than relaxing to a continuum and projecting back, is what returns the loss\-versus\-wall\-clock frontier that practitioners want\.

### D\.3Optimal Resource Allocation

MOSAIC solves the allocation problem numerically by enumerating realizable architectures and searching over their feasible execution plans\. This section asks what an additional closed\-form prescription in the style ofhoffmann2022trainingwould look like, and what assumptions would be required to state one\. Nothing below replaces the discrete search of Section[7](https://arxiv.org/html/2608.10605#S7)\.

#### D\.3\.1Reduction to aCpeakC\_\{\\text\{peak\}\}Law

The exogenous hardware\-time budget is the raw ceilingCpeak=Fpeak​Ndev​TtrainC\_\{\\text\{peak\}\}=F\_\{\\text\{peak\}\}N\_\{\\text\{dev\}\}T\_\{\\text\{train\}\}of Section[2\.1](https://arxiv.org/html/2608.10605#S2.SS1)together with the goodput fractionηgood​\(Ndev\)\\eta\_\{\\text\{good\}\}\(N\_\{\\text\{dev\}\}\)\. Both are fixed before an architecture is selected, soCpeak​ηgoodC\_\{\\text\{peak\}\}\\eta\_\{\\text\{good\}\}is the counterpart of the compute budget in a Chinchilla\-style prescription, whereas the deliverable budgetCdeliverableC\_\{\\text\{deliverable\}\}is not, since it also carries the best attainableMFU⋆⁡\(Z\)=maxPℓ∈𝒫feas​\(Z,B\)⁡MFU⁡\(Z,Pℓ,B\)\\operatorname\{MFU\}^\{\\star\}\(Z\)=\\max\_\{P\_\{\\ell\}\\in\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)\}\\operatorname\{MFU\}\(Z,P\_\{\\ell\},B\)over the feasible plan set𝒫feas​\(Z,B\)\\mathcal\{P\}\_\{\\text\{feas\}\}\(Z,B\)of Section[2\.2](https://arxiv.org/html/2608.10605#S2.SS2)\. Loss decreases monotonically in the token count, so saturating Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\) recovers Property 3,

D⋆​\(Z,Cpeak\)=Cpeak​MFU⋆⁡\(Z\)​ηgood6​Nact​\(Z\)\.D^\{\\star\}\(Z,C\_\{\\text\{peak\}\}\)=\\frac\{C\_\{\\text\{peak\}\}\\,\\operatorname\{MFU\}^\{\\star\}\(Z\)\\,\\eta\_\{\\text\{good\}\}\}\{6N\_\{\\text\{act\}\}\(Z\)\}\.\(26\)SinceMFU⋆⁡\(Z\)\\operatorname\{MFU\}^\{\\star\}\(Z\)does not depend on the training\-window length, substituting Equation \([26](https://arxiv.org/html/2608.10605#A4.E26)\) into Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) collapses the law to

ℒ⋆​\(Z,Cpeak\)=ℒ∞​\(Z\)\+Ψ​\(Z\)Cpeakβ,Ψ​\(Z\)=b​\(6​Nact​\(Z\)MFU⋆⁡\(Z\)​ηgood\)β\\mathcal\{L\}^\{\\star\}\(Z,C\_\{\\text\{peak\}\}\)=\\mathcal\{L\}\_\{\\infty\}\(Z\)\+\\frac\{\\Psi\(Z\)\}\{C\_\{\\text\{peak\}\}^\{\\beta\}\},\\qquad\\Psi\(Z\)=b\\left\(\\frac\{6N\_\{\\text\{act\}\}\(Z\)\}\{\\operatorname\{MFU\}^\{\\star\}\(Z\)\\,\\eta\_\{\\text\{good\}\}\}\\right\)^\{\\\!\\beta\}\(27\)whereℒ∞​\(Z\)\\mathcal\{L\}\_\{\\infty\}\(Z\)collects the three terms of Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) that do not involveDD, together with the irreducible constantee\. We note thatD⋆∝CpeakD^\{\\star\}\\propto C\_\{\\text\{peak\}\}at any fixed architecture, provided no external data or tokens\-per\-parameter cap binds, though the globally optimal token count doesn’t need to be linear inCpeakC\_\{\\text\{peak\}\}because the selected architecture can change the constantMFU⋆⁡\(Z\)/Nact​\(Z\)\\operatorname\{MFU\}^\{\\star\}\(Z\)/N\_\{\\text\{act\}\}\(Z\)\.

#### D\.3\.2Conditions for an allocation law

The reduction above is exact for each discrete candidate, but a continuous optimal solution needs more assumptions, and what follows is a local, regime\-specific approximation rather than a universal law\. We writer≡1−Sr\\equiv 1\-SandN≡NtotN\\equiv N\_\{\\text\{tot\}\}, and we assume \(1\) an interior memory\-slack region where any memory or geometry boundary constraint is non\-existent, \(2\) the same execution layout stays optimal locally, and \(3\)MFU⋆⁡\(N,r,G\)\\operatorname\{MFU\}^\{\\star\}\(N,r,G\)is differentiable\. We define the local MFU log\-slopesωN=∂log⁡MFU⋆/∂log⁡N\\omega\_\{N\}=\\partial\\log\\operatorname\{MFU\}^\{\\star\}/\\partial\\log N,ωS=∂log⁡MFU⋆/∂log⁡r\\omega\_\{S\}=\\partial\\log\\operatorname\{MFU\}^\{\\star\}/\\partial\\log r, andωG=−∂log⁡MFU⋆/∂log⁡G\\omega\_\{G\}=\-\\,\\partial\\log\\operatorname\{MFU\}^\{\\star\}/\\partial\\log G, so thatωG\>0\\omega\_\{G\}\>0means MFU falls as experts become finer\. After writing the four terms of Equation \([10](https://arxiv.org/html/2608.10605#S4.E10)\) asℒcap=a​N−α\\mathcal\{L\}\_\{\\text\{cap\}\}=aN^\{\-\\alpha\},ℒdata=b​D−β\\mathcal\{L\}\_\{\\text\{data\}\}=bD^\{\-\\beta\},ℒsp=c​r−λ\\mathcal\{L\}\_\{\\text\{sp\}\}=cr^\{\-\\lambda\}andℒint=j​N−γ​r−δ​G−η\\mathcal\{L\}\_\{\\text\{int\}\}=jN^\{\-\\gamma\}r^\{\-\\delta\}G^\{\-\\eta\}, and eliminatingDDthrough Equation \([26](https://arxiv.org/html/2608.10605#A4.E26)\), stationarity in\(log⁡N,log⁡r,log⁡G\)\(\\log N,\\log r,\\log G\), we obtain

α​ℒcap\\displaystyle\\alpha\\,\\mathcal\{L\}\_\{\\text\{cap\}\}=β​ℒdata​\(1−ωN−γ​ωG/η\),\\displaystyle=\\beta\\,\\mathcal\{L\}\_\{\\text\{data\}\}\\bigl\(1\-\\omega\_\{N\}\-\\gamma\\omega\_\{G\}/\\eta\\bigr\),λ​ℒsp\\displaystyle\\lambda\\,\\mathcal\{L\}\_\{\\text\{sp\}\}=β​ℒdata​\(1−ωS−δ​ωG/η\),η​ℒint=β​ωG​ℒdata\.\\displaystyle=\\beta\\,\\mathcal\{L\}\_\{\\text\{data\}\}\\bigl\(1\-\\omega\_\{S\}\-\\delta\\omega\_\{G\}/\\eta\\bigr\),\\qquad\\eta\\,\\mathcal\{L\}\_\{\\text\{int\}\}=\\beta\\,\\omega\_\{G\}\\,\\mathcal\{L\}\_\{\\text\{data\}\}\.\(28\)Since all four loss components and their exponents are positive, these stationarity equations require necessary sign conditions for an interior solution\. In particular, theGGequation requiresωG\>0\\omega\_\{G\}\>0, while theNNandrrequations require1−ωN−γ​ωG/η\>01\-\\omega\_\{N\}\-\\gamma\\omega\_\{G\}/\\eta\>0and1−ωS−δ​ωG/η\>01\-\\omega\_\{S\}\-\\delta\\omega\_\{G\}/\\eta\>0, respectively\. Without a systems penalty on finer experts,ωG=0\\omega\_\{G\}=0, the third stationarity condition cannot hold\. The objective is then strictly decreasing inGGwithin this regime, so no interiorGGoptimum exists and the solution lies at the largest admissible expert split factor\. Where memory binds, a geometry boundary is active, or the optimal layout changes, the corresponding KKT or nonsmooth terms must be included instead\.

Note that Equation \([28](https://arxiv.org/html/2608.10605#A4.E28)\) alone does not give a power law, since the log\-slopes may drift as the architecture moves\. The key additional assumption is that they stay approximately constant within one systems regime, so thatMFU⋆≃m0​\(N/N0\)ωN​\(r/r0\)ωS​\(G/G0\)−ωG\\operatorname\{MFU\}^\{\\star\}\\simeq m\_\{0\}\(N/N\_\{0\}\)^\{\\omega\_\{N\}\}\(r/r\_\{0\}\)^\{\\omega\_\{S\}\}\(G/G\_\{0\}\)^\{\-\\omega\_\{G\}\}\. Under this approximation the four terms stay in fixed ratios and the optimum traces the regime\-local path

r⋆∝\(N⋆\)α/λ,D⋆∝\(N⋆\)α/β,G⋆∝\(N⋆\)\(α−γ−α​δ/λ\)/η\\boxed\{\\;r^\{\\star\}\\propto\\left\(N^\{\\star\}\\right\)^\{\\alpha/\\lambda\},\\quad D^\{\\star\}\\propto\\left\(N^\{\\star\}\\right\)^\{\\alpha/\\beta\},\\quad G^\{\\star\}\\propto\\left\(N^\{\\star\}\\right\)^\{\(\\alpha\-\\gamma\-\\alpha\\delta/\\lambda\)/\\eta\}\\;\}\(29\)Writingp=α/λp=\\alpha/\\lambda,ϱ=α/β\\varrho=\\alpha/\\betaandg=\(α−γ−α​δ/λ\)/ηg=\(\\alpha\-\\gamma\-\\alpha\\delta/\\lambda\)/\\etafor these exponents, andν=\[\(1−ωN\)\+\(1−ωS\)​p\+ϱ\+ωG​g\]−1\\nu=\[\(1\-\\omega\_\{N\}\)\+\(1\-\\omega\_\{S\}\)p\+\\varrho\+\\omega\_\{G\}g\]^\{\-1\}, the same path in terms of the purchased budget is

N⋆∝Cpeakν,1−S⋆∝Cpeakp​ν,D⋆∝Cpeakϱ​ν,G⋆∝Cpeakg​ν\.\\boxed\{\\;N^\{\\star\}\\propto C\_\{\\text\{peak\}\}^\{\\nu\},\\quad 1\-S^\{\\star\}\\propto C\_\{\\text\{peak\}\}^\{p\\nu\},\\quad D^\{\\star\}\\propto C\_\{\\text\{peak\}\}^\{\\varrho\\nu\},\\quad G^\{\\star\}\\propto C\_\{\\text\{peak\}\}^\{g\\nu\}\.\\;\}\(30\)These relations hold conditional on fixed hardware, device count, batch envelope, goodput model, and systems regime\. Layout transitions, memory boundaries, and topology changes alter the log\-slopes and therefore produce different allocation phases\.

## Appendix EPerformance Model Input/Outputs

The listings below describe ScalePlan \([https://github\.com/dmlc/ScalePlan](https://github.com/dmlc/ScalePlan)\), the performance model released with this paper\. Beyond scoring one plan, it enumerates the execution plansPℓ=\(PTP,PEP,PCP,PPP,PDP,Bmicro,Ackpt\)P\_\{\\ell\}=\(P\_\{\\text\{TP\}\},P\_\{\\text\{EP\}\},P\_\{\\text\{CP\}\},P\_\{\\text\{PP\}\},P\_\{\\text\{DP\}\},B\_\{\\text\{micro\}\},A\_\{\\text\{ckpt\}\}\)for a given device count, discards those that breach the memory cap, and returns the survivors ranked by predicted MFU\. That search is what supplies the inner maximization over𝒫feas\\mathcal\{P\}\_\{\\text\{feas\}\}of Equation \([C\-FLOPs](https://arxiv.org/html/2608.10605#S2.Ex1)\), so MOSAIC calls it once per candidate architecture rather than assuming a layout\. The field names below are those of the tool, and the trailing comment on each line gives the corresponding paper symbol where one exists\.

1model:

2n\_layers:48

3hidden\_sz:6144

4inter\_sz:15360

5n\_q\_heads:48

6n\_kv\_heads:8

7head\_dim:128

8vocab\_sz:257536

9precision:"bf16"

10sdpa\_precision:"bf16"

11glu:true

12rotary\_embeds:true

13dropout:false

14tie\_embeddings:false

15

16moe:

17n\_experts:128

18experts\_per\_token:3

19capacity\_factor:1\.0

20expert\_inter\_sz:3840

21moe\_frequency:1

22expert\_tp\_degree:1

23

24search:

25num\_devices:512

26

27performance:

28activation\_checkpointing\_type:full

29

30optimizer:

31optimizer\_type:muon

32

33data:

34gbs:8192

35seqlen:8192

36microbatch\_sz:2

37

38hardware:

39node\_type:p6\-b200\.48xlarge

Listing 1:Performance model input specification\.Output example, where the layout degrees are\(PTP,PEP,PPP,PCP,PDP\)\(P\_\{\\text\{TP\}\},P\_\{\\text\{EP\}\},P\_\{\\text\{PP\}\},P\_\{\\text\{CP\}\},P\_\{\\text\{DP\}\}\)ofPℓP\_\{\\ell\}and the reported MFU and memory are theMFU​\(Z,Pℓ,B\)\\mathrm\{MFU\}\(Z,P\_\{\\ell\},B\)andMem​\(Z,Pℓ,B\)\\mathrm\{Mem\}\(Z,P\_\{\\ell\},B\)predictions of Section[5](https://arxiv.org/html/2608.10605#S5)\.

1Rank1:

2Configuration:tp=1,ep=8,pp=16,cp=1,dp=32

3AnalyticalMFU:13\.69%

4Memoryperdevice:21\.30GB

5Iterationtime:44\.892s

Listing 2:Performance model output\.
## Appendix FResults

### F\.1Geometry Scale Ladders

Table[9](https://arxiv.org/html/2608.10605#A6.T9)lists the rungs of the geometry scaling ladder of Section[7\.1](https://arxiv.org/html/2608.10605#S7.SS1)that our search enumerates\. Each row is one seedqq, giving the backbone dimensions the seed fixes and, at a representative MoE configuration, the parameter counts and sparsity that backbone realizes\.

qqLlayersL\_\{\\text\{layers\}\}nheadn\_\{\\text\{head\}\}dddffd\_\{\\text\{ff\}\}NtotN\_\{\\text\{tot\}\}\(B\)∗NactN\_\{\\text\{act\}\}\(B\)∗SS416162048512015\.50\.70\.954624243072768053\.22\.30\.95783232409610240127\.35\.40\.958104040512012800250\.010\.40\.958124848614415360433\.717\.90\.959145656716817920690\.528\.20\.9591664648192204801032\.841\.90\.9591872729216230401472\.859\.50\.960Table 9:Seeds of the model\-scale ladder used in our experiments\. The geometry seedqqsets the backbone depthLlayers=4​qL\_\{\\text\{layers\}\}=4q, head countnhead=4​qn\_\{\\text\{head\}\}=4q, residual widthd=128​nheadd=128\\,n\_\{\\text\{head\}\}, and reference dense FFN widthdff=2\.5​dd\_\{\\text\{ff\}\}=2\.5\\,d\. The last three columns give the total and active parameter counts and the induced sparsitySSat an example MoE configurationE=128E\\\!=\\\!128,K=3K\\\!=\\\!3,G=4G\\\!=\\\!4\. The MoE axes\(E,K,G\)\(E,K,G\)are searched at each rung\.∗Non\-embedding parameter counts\.
### F\.2Implications for Sparse\-MoE Parallelism Design

This appendix develops the design rules summarized in Section[7\.2\.2](https://arxiv.org/html/2608.10605#S7.SS2.SSS2)\. Reading the measured best MFU achieved by different parallelism configurations in the validation sweeps yields two design rules and one caveat that bounds them\.

- •The dominant lever is not any single degree but the pipeline fill ratiorfill=nmb/PPPr\_\{\\text\{fill\}\}=n\_\{\\text\{mb\}\}/P\_\{\\text\{PP\}\}, the microbatches in flight per pipeline stage\.Within a sweep, layout choice moves measured MFU by up to3\.1×3\.1\\times\(18B\) and2\.1×2\.1\\times\(700M\), and almost all of that spread tracksrfillr\_\{\\text\{fill\}\}, since 1F1B leaves each stage idle for a bubble fraction of\(PPP−1\)/\(PVPP​nmb\)\(P\_\{\\text\{PP\}\}\{\-\}1\)/\(P\_\{\\text\{VPP\}\}\\,n\_\{\\text\{mb\}\}\), so throughput collapses asrfillr\_\{\\text\{fill\}\}falls toward11and saturates oncerfillr\_\{\\text\{fill\}\}is large\. The 18B pipeline sweep makes this explicit, where atPPP=16P\_\{\\text\{PP\}\}\{=\}16, MFU climbs asrfillr\_\{\\text\{fill\}\}goes0\.25→0\.5→10\.25\\to 0\.5\\to 1, then plateaus at∼10\\sim\\\!10–12%12\\%forrfill≥4r\_\{\\text\{fill\}\}\\geq 4\. This reframes the two parallelism knobs that people tune independently, since deeper pipelines and larger microbatches both*shrink*rfillr\_\{\\text\{fill\}\}and are only free once the fill ratio is already comfortable\. Concretely, every extra pipeline stage or microbatch that dropsrfillr\_\{\\text\{fill\}\}below∼4\\sim\\\!4costs measurable MFU, whereas the same change atrfill≥4r\_\{\\text\{fill\}\}\\geq 4is nearly free\.
- •Keep the expert all\-to\-all inside the scale\-up \(NVLink\) domain, and, once the fill ratio is satisfied, pipeline depth is what should absorb the rest of the budget\.Expert parallelism helps up to the node boundary and hurts past it, where MFU rises monotonically withPEPP\_\{\\text\{EP\}\}in the 700M sweep \(forPEP∈\{1,2,4,8\}P\_\{\\text\{EP\}\}\\in\\\{1,2,4,8\\\}at fixed depth\) and drops fromPEP=8P\_\{\\text\{EP\}\}\{=\}8toPEP=16P\_\{\\text\{EP\}\}\{=\}16in both 5\.4B sweeps as the group leaves the 8\-GPU node\. OncePEPP\_\{\\text\{EP\}\}is pinned to the domain, moderate pipeline depth is the cheapest remaining axis provided it does not pushrfillr\_\{\\text\{fill\}\}below the plateau, and the 5\.4B optima usePPP∈\{2,4\}P\_\{\\text\{PP\}\}\{\\in\}\\\{2,4\\\}because at these depthsrfillr\_\{\\text\{fill\}\}stays≥8\\geq 8\. This is the actionable ordering the model exposes, where we setPEPP\_\{\\text\{EP\}\}to the scale\-up domain, then add pipeline depth only whilerfillr\_\{\\text\{fill\}\}remains large\.
- •Spend the FFN\-parallelism budget on expert parallelism rather than expert tensor parallelism\.For the expert FFN this is the expected ordering, since expert parallelism partitions the experts with one all\-to\-all per layer, whereas tensor parallelism adds two all\-reduces per layer on the critical path\.

Scoped this way, the rules collapse the viable search region sharply and in a model\-checkable order, where we fixPTP=1P\_\{\\text\{TP\}\}\{=\}1, setPEPP\_\{\\text\{EP\}\}to the scale\-up domain, then choose the\(PPP,Bmicro\)\(P\_\{\\text\{PP\}\},B\_\{\\text\{micro\}\}\)pair that keeps the pipeline fill ratiorfill≥4r\_\{\\text\{fill\}\}\\geq 4at the target global batch\. This retains the measured optimum in every sweep that is not memory\-bound while discarding the configurations the analytical model also ranks worst, which lowers the cost of the hardware\-confirmation step\. Long\-context, memory\-bound, or wider\-NVLink regimes shift the thresholds along the principles above\.

### F\.3The Training Stack Behind the Staged Runs

The staged runs of Section[7\.5](https://arxiv.org/html/2608.10605#S7.SS5)are launched on the same discrete grid that Stage\-2 enumerates, but their realized throughput is set by a production training stack whose overheads the analytical performance model only approximates\. Each run is driven by a PyTorch Lightning, NeMo, and Megatron\-Core pipeline with one process per GPU and eight GPUs per p6\-b200 node\. We use the Megatron\-Core MoE training stack777https://github\.com/NVIDIA/Megatron\-LM/tree/main/megatron/core/transformer/moeas is, both for the performance\-model validation and for the staged runs, so we measure a standard stack and not a custom one\. From that stack we enable the grouped\-GEMM expert path \(so the per\-expert matrix multiplications for a layer are issued as one batched kernel\), cross\-entropy loss fusion, the all\-to\-all token dispatcher for expert parallelism, and shared\-expert overlap \(which hides the shared\-expert compute behind the routed all\-to\-all\)\. Fine\-grained activation recomputation is treated as a sweep axis inside the MFU search rather than fixed, since it trades memory for recompute cost and thereby moves the feasible\-layout boundary\. We deliberately leave the rest of the stack’s kernels and fusions off, and in particular we do not use parallel folding or DeepEP, and we leave many of the newer optimizations in the stack as future work\. An important point is that our performance model is itself an evolving artifact that mirrors one systems stack at a given point in time, so as we turn on new kernels or change the stack, the model is simply re\-calibrated or replaced by an updated one rather than being a fixed assumption of the framework\.

Several loop\-level effects open a gap between this realized MFU and the value the analytical model predicts, and they are the reason the analytical model at times is systematically optimistic \(Section[7\.5](https://arxiv.org/html/2608.10605#S7.SS5)\)\. Pipeline parallelism over large number of micro\-batches leaves an idle bubble that a steady\-state FLOPs count does not see\. Full activation recomputation, used to fit the model in memory, adds an extra forward pass per layer inside the backward, raising step time without adding any of the idealized compute that MFU credits\. The expert all\-to\-all dispatch is exposed collective time whose cost depends on whether an expert\-parallel group stays inside a node or crosses the fabric, and the fp32 router is a precision island inside an otherwise bf16 model\. Together with asynchronous S3 checkpointing, periodic garbage collection, and launch jitter, these are wall\-clock the device\-side model does not fully charge, so it over\-predicts MFU while preserving the ranking the selection relies on\.

A second class of overhead lives entirely on the host\. Our tokenizer is a fast Hugging Face tokenizer, downloaded once per node from S3 into shared memory and patched to be thread\-safe because the loader calls it from worker threads\. The collator pads every micro\-batch to a multiple of eight and then to the full20482048\-token context, since pipeline parallelism requires uniform shapes across stages\. The streaming dataloader shards only along the data\-parallel axis, replicates identical batches across the tensor\-, pipeline\-, and expert\-parallel ranks of a group\. Because we run a single loader worker per rank with a prefetch depth of two, host\-side tokenization, collation, and the host\-to\-device copy can stall the accelerators\. Such input\-pipeline stalls lower realized MFU and are invisible to a device\-side performance model, which is one more reason in the gap between the measured throughput and the prediction\.

Finally, the realization is specific to the accelerator and its interconnect, so the same architecture can land at a different optimum on different hardware\. A p6\-b200 node888https://aws\.amazon\.com/ec2/instance\-types/p6/pairs eight NVIDIA B200 GPUs with up to14\.414\.4TB/s of bidirectional intra\-node NVLink and a fourth\-generation Elastic Fabric Adapter delivering3\.23\.2Tbps \(400400Gbps per GPU\) between nodes\. On an H200 platform the per\-device peak FLOPs, the HBM capacity, and the NVLink bandwidth are all lower, which shifts both the memory cap and the throughput\-optimal layout\. Because intra\-node collectives ride NVLink and NVSwitch while inter\-node collectives ride EFA RDMA, whether the expert all\-to\-all is confined to a node or spills across the fabric can swing MFU sharply\. The optimum MOSAIC returns is therefore tied to the device peak FLOPs and the bandwidth constants of the target platform, and it must be resolved, not merely rescaled, when the hardware changes\.