Paper deep dive
Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts
Soumajyoti Sarkar, Yuxin Tang, Sheng Zha
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 8/16/2026, 2:56:18 AM
Summary
The paper introduces MOSAIC, a framework for systems-aware scaling of sparse Mixture-of-Experts (MoE) language models. It argues that traditional compute-optimal scaling laws, which optimize for model FLOPs, fail to account for hardware efficiency constraints like memory footprint and communication costs. MOSAIC couples a predictive scaling law with a calibrated performance model to estimate Model FLOPs Utilization (MFU) and other systems metrics. By formulating architecture and systems co-design as a joint optimization problem, MOSAIC identifies that optimal sparsity emerges from cluster-specific systems constraints rather than pure compute budgets, leading to more efficient frontier model training.
Entities (6)
Relation Signals (5)
MOSAIC → developedby → Amazon AGI Foundations
confidence 99% · Authors from Amazon AGI Foundations develop MOSAIC
MOSAIC → optimizes → Sparse Mixture-of-Experts
confidence 95% · We instantiate the framework for sparse Mixture-of-Experts (MoE) language models
MOSAIC → estimates → Model FLOPs Utilization
confidence 92% · MOSAIC couples a predictive scaling law with a calibrated performance model that estimates Model FLOPs Utilization (MFU)
Sparse Mixture-of-Experts → uses → Scaling Law
confidence 88% · We fit a scaling law on sparse MoE models
MOSAIC → validateson → NVIDIA B200
confidence 85% · We search the architecture grid on a cluster of NVIDIA B200 GPUs
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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 from $104$ million to $2.7$ billion and total model sizes reaching $79$ billion 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.
Tags
Links
- Source: https://arxiv.org/abs/2608.10605v1
- Canonical: https://arxiv.org/abs/2608.10605v1
Trouble viewing inline? Open PDF directly →
Full Text
149,456 characters extracted from source content.
Expand or collapse full text
Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts Sarkar @amazon.com AGI Foundations Tang @amazon.com AGI Foundations Zha @amazon.com 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 from 104104 million to 2.72.7 billion and total model sizes reaching 7979 billion 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 1 Introduction 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 size N) and data (token count D) 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, hoffmann2022training devise a joint scaling law for loss in the model size N and the number of training tokens D, 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 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, anthony2024codesign show 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 of E experts and routing each token to a subset of them. This decouples the active parameters NactN_act that set per-token model-FLOPs from the total stored parameters NtotN_tot, with the expert count E as 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, including abnar2025parameters, who show that the sparsity factor S, roughly the proportion of NtotN_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 and Figure 1 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. ludziejewski2025joint fit 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 optimal E is where the loss surface meets a memory-aware budget rather than the E that minimizes loss at fixed NactN_act. wan2026holistic push 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. 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 in sun2026hardware which 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). (a) Model FLOPs selects a boundary MoE sparsity. Under a fixed Cmodel=3.3×1021C_model=3.3× 10^21, the fitted loss decreases over the calibrated sparsity range, so S⋆=SmaxS =S_ (left). We do not cap any model dimension, and at each sparsity till SmaxS_ (0.985 here), the expert count and activated parameters are optimized under the budget. Under a hardware budget CpeakC_peak with 4 AWS p6-B200 nodes for 5 days, this prescription is infeasible because configurations beyond the feasible boundary of S=0.96 deliver fewer CdeliverableC_deliverable FLOPs (red region) than what the fixed prescription S⋆S requires (right). (b) MOSAIC yields an interior feasible sparsity under the hardware budget. The hardware-aware loss ℒhw⋆(S)L _hw(S) induced by CdeliverableC_deliverable is minimized at the interior sparsity Scluster⋆ [rgb]0,0.55,0 [named]pgfstrokecolorrgb0,0.55,0S _cluster and is 0.031 nats lower than at S⋆=Smax [rgb]1,0,0 [named]pgfstrokecolorrgb1,0,0S =S_ . 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 under 15%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 to 250250B 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. 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. We then describe the components that instantiate MOSAIC with the sparse MoE architecture: (1) sparse MoE scaling laws in Section 4 and, (2) the performance model in Section 5. We make the application of MOSAIC for sparse MoEs concrete in Section 6, 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. We provide a snapshot of the most shared symbols used in the rest of the paper in Table 1. Symbol Meaning Architecture LlayersL_layers Number of transformer blocks. d Residual-stream width (dmodeld_model). dffd_f Hidden width of the dense FFN that a sparse MoE block replaces. dexpertd_expert Hidden width of a single expert. E Total number of experts per sparse MoE block. K Top-K active experts per token. G expert split factor, G=dff/dexpertG=d_f/d_expert. Figure 14 in Appendix B visualizes the construction. NactN_act Active parameters per token (FLOPs-relevant). NtotN_tot Total stored parameters of the model (model-capacity). S Sparsity, S=1−Nact/NtotS=1-N_act/N_tot (fraction of parameters not activated per token). Routing and load balancing H∈ℝT×dH ^T× d Pre-MoE hidden states for a batch of T tokens (Algorithm 1). ϕ∈ℝEφ ^E Per-expert routing bias (LFLB controller). μb _b Bias-update rate of the LFLB controller. ρ Proposal-to-selection ratio for Gumbel routing, giving proposal size M=min(ρK,E)M= (ρ K,E). si′(x)s _i(x), si(x)s_i(x) Clean sigmoid router score and the selection score si(x)=si′(x)+ϕis_i(x)=s _i(x)+ _i that adds the balancing bias (Section 3.3). K(x)T_K(x) Set of top-K experts selected for token x. MOSAIC decision variables Z Architecture decision variable of MOSAIC, a realizable sparse MoE geometry with induced scaling law dimensions (Ntot,S,G)(N_tot,S,G) (Section 2). PℓP_ Execution plan, Pℓ=(PTP,PEP,PCP,P,PDP,Bmicro,Ackpt)P_ =(P_TP,P_EP,P_CP,P_P,P_DP,B_micro,A_ckpt), giving the parallelism degrees along tensor (TP), expert (EP), context (CP), pipeline (P), and data parallel (DP) groups together with the micro-batch size and the activation-checkpointing mode. B Batch envelope, B=(Bglobal,Tseq)B=(B_global,T_seq), giving nmb=Bglobal/(PDPBmicro)n_mb=B_global/(P_DP\,B_micro) microbatches per optimizer step. D Number of training tokens. CmodelC_model Model FLOPs, Cmodel=6N(Z)DC_model=6\,N(Z)\,D, where the activated parameters N(Z)N(Z) are NactN_act for sparse MoEs. TtrainT_train Wall-clock training window on the cluster. CpeakC_peak Raw hardware ceiling, Cpeak=FpeakNdevTtrainC_peak=F_peak\,N_dev\,T_train, where NdevN_dev is the device count. FpeakF_peak Per-device peak throughput. CdeliverableC_deliverable Deliverable model FLOPs, Cdeliverable(Z,Pℓ,B)=Cpeak⋅MFU(Z,Pℓ,B)⋅ηgood(Ndev)C_deliverable(Z,P_ ,B)=C_peak·MFU(Z,P_ ,B)· _good(N_dev), MFU is defined in Equation (1) and ηgood _good is the goodput fraction. Table 1: Symbols used in the paper. 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. 2 MOSAIC: Model Optimization via Systems-Aware TraIning Co-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, we define the resulting architecture-dependent deliverable budget, we formulate the joint optimization in Section 2.2, and we instantiate it for sparse MoEs in Section 2.3. Figure 2 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.1 From Hardware Budget to Deliverable Model FLOPs A model FLOPs budget treats every allocation with the same Cmodel(Z,D)=6N(Z)DC_model(Z,D)=6\,N(Z)\,D as equally expensive, where N(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 of NdevN_dev devices, each with peak per-device throughput FpeakF_peak, provides the raw hardware ceiling Cpeak=FpeakNdevTtrainC_peak=F_peak\,N_dev\,T_train over a training window TtrainT_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 plan PℓP_ specifies how the model is sharded and its computation scheduled across the GPU fleet (meta2025llama3isca). For architecture Z, execution plan PℓP_ , and batch envelope B, let titer(Z,Pℓ,B)t_iter(Z,P_ ,B) denote the wall-clock time of one optimizer step processing BglobalTseqB_global\,T_seq tokens, where TseqT_seq is the sequence length. We define MFU as MFU(Z,Pℓ,B)=6N(Z)BglobalTseqtiter(Z,Pℓ,B)NdevFpeakMFU(Z,P_ ,B)\;=\; 6\,N(Z)\;B_global\,T_seqt_iter(Z,P_ ,B)\;N_dev\;F_peak (1) It is the model FLOPs per token executed per second, normalized by the peak FLOPs rate FpeakNdevF_peakN_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) _good(N_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_deliverable(Z,P_ ,B)\;=\;C_peak\;·\;MFU(Z,P_ ,B)\;·\; _good(N_dev). (2) Unlike a conventional compute budget, CdeliverableC_deliverable depends on the architecture and the layout PℓP_ used to execute it. A candidate training configuration (Z,D)(Z,D) is feasible to train on the cluster within the window TtrainT_train only if it satisfies the constraint Cmodel(Z,D)≤maxPℓ∈feas(Z,B)Cdeliverable(Z,Pℓ,B)C_model(Z,D)\;≤\; _P_ _feas(Z,B)\,C_deliverable(Z,P_ ,B) (3) where feas(Z,B)P_feas(Z,B) collects the execution plans that fit the cluster constraint, defined in Section 2.2. 2.2 Joint 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 architecture Z, drawn from a finite set discZ_disc of realizable architectures, each carrying the scaling law dimensions its geometry induces, the training tokens D, and the execution plan PℓP_ of Section 2.1. The hardware inputs are (Fpeak,Ndev,Ttrain,Mdev)(F_peak,N_dev,T_train,M_dev), where MdevM_dev is the per-device memory cap, and the batch envelope B=(Bglobal,Tseq)B=(B_global,T_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)≤MdevP_feas(Z,B)\;=\; \\,P_ _disc(N_dev)\;:\;Mem(Z,P_ ,B)≤ M_dev\, \ (4) where disc(Ndev)P_disc(N_dev) is the finite set of realizable execution plans the target training stack exposes on NdevN_dev devices, comprising the parallelism degrees that tile the cluster, the micro-batch size BmicroB_micro, and the activation-checkpointing mode AckptA_ckpt. Over this feasible plan set, we solve the following optimization minZ∈disc,D _Z _disc,\;D ℒ(Z,D) (Z,D ) (5) subject to Cmodel(Z,D)≤maxPℓ∈feas(Z,B)Cdeliverable(Z,Pℓ,B). C_model(Z,D)\;≤\; _P_ _feas(Z,B)\,C_deliverable(Z,P_ ,B). (C-FLOPs) The objective ℒL is the scaling law, which predicts loss from the architecture specific dimension Z and the data D. The loss functional form evaluates Z through the scaling law dimensions it induces, whereas MFU(.)MFU(.) and Mem(.)Mem(.) evaluate the concrete architecture itself, so all three are well-defined functions of Z even 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)P_feas(Z,B), which means that the maximum never returns an infeasible plan. Equation (C-FLOPs) is the feasibility condition of Equation (3). Equivalently, MOSAIC chooses for each architecture the best feasible plan as Pℓ⋆(Z)=argmaxPℓ∈feas(Z,B)Cdeliverable(Z,Pℓ,B).P_ (Z)\;=\; _P_ _feas(Z,B)\;C_deliverable(Z,P_ ,B). (6) We note that Equation (5) is a bilevel non-convex mixed-integer nonlinear program, an optimization problem that is bilevel because of the embedded maximization over feasP_feas in Equation (C-FLOPs). It is mixed-integer nonlinear since the layout degrees are integer and the architecture is drawn from the discrete set discZ_disc. The constraint functions MFU(⋅)MFU(·), Mem(⋅)Mem(·) are nonlinear, and non-convex because of the discrete variables of the realizable plan set disc(Ndev)P_disc(N_dev) together with the non-convexity of MFU and memory in PℓP_ . We collect in Appendix D.2 three structural properties of this program, including a closed-form elimination of D, that make its exact solution by enumeration tractable. 2.3 Instantiating MOSAIC for Sparse MoEs We put Equation (5) to be architecture-agnostic. Applying it requires three components, a realizable architecture space induced by Z, a predictive loss function ℒ(Z,D)L(Z,D), and predictors for MFU(Z,Pℓ,B)MFU(Z,P_ ,B) and Mem(Z,Pℓ,B)Mem(Z,P_ ,B). We define the sparse MoE realization of Z in Section 3, under which discZ_disc becomes the finite set of realizable sparse MoE geometries and the activated parameter count becomes N(Z)=NactN(Z)=N_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)P_disc(N_dev) for our MoE models in Appendix D.1. We discuss each of these three components applied to sparse MoEs in the following sections. Figure 1 instantiates the constraint of Equation 3 with real numbers where a Cmodel=3.3×1021C_model=3.3× 10^21 prescription on 4 p6-B200 nodes for 5 days becomes infeasible beyond Sfeasmax=0.96S_feas =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. 3 Sparse MoE Pretraining: Scaling Recipe In this section, we discuss the dimensions Z for 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 of krajewski2024scaling and abnar2025parameters with some adjustments, and we describe those differences in the following sections. 3.1 Training 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 the K routed 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. 3.2 MoE Layer We use LlayersL_layers to denote the number of transformer blocks, d≡dembdd\!≡\!d_embd as the residual stream width, and dffd_f as the hidden width of the corresponding dense feed-forward layer that a sparse MoE layer replicates into E copies 333dffd_f is conventionally a fixed multiple of the residual width d. For example, GPT-2 (radford2019language) uses dff=4d_f=4d.. Each MoE layer consists of E experts fii=1E\f_i\_i=1^E, where every expert is itself a two-layer feed-forward network of the form fi(x)=Wi(2)σ(Wi(1)x)f_i(x)=W^(2)_i\,σ\! (W^(1)_ix ) with Wi(1)∈ℝdexpert×dW^(1)_i ^d_expert× d, Wi(2)∈ℝd×dexpertW^(2)_i ^d× d_expert, and σ(⋅)σ(·) a gated activation, in our runs being SwiGLU. We keep dffd_f and dexpertd_expert as separate symbols because a sparse MoE layer can split the reference FFN of width dffd_f into narrower experts of width dexpertd_expert, so the two coincide only when each expert is a full-width dense FFN, something we will explore in Section 4. We follow the Switch-style token-choice routing formulation (fedus2022switch; liu2023sparsemixer). A learnable router Wr∈ℝE×dW_r ^E× d produces a per-token score vector (x)∈ℝE s(x) ^E, and the top-K experts under (x) 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)=TopK((x))y(x)\;=\; _i\,∈\,T_K(x) π_i(x)\,f_i(x), _K(x)\;=\;Top\!K ( s(x) ) (7) where π~i(x) π_i(x) are the probabilities the selected set K(x)T_K(x). In the following section, we describe the routing operations that result in the TopK(.)Top\!K(.) outputs and the associated probabilities π~i(x) π_i(x) and we treat Equation 7 as the canonical sparse MoE layer definition. Algorithm 1 Sparse MoE token-choice routing for a single MoE layer, batch of T tokens routed to E experts 0: Pre-MoE hidden states H∈ℝT×dH ^T× d (attention output of the current block added to the previous-layer residual stream); router weights Wr∈ℝE×dW_r ^E× d; per-expert bias ϕ∈ℝEφ ^E; hyperparameters K (top-K), Gumbel proposal size M=min(ρK,E)M= (ρ K,E), Gumbel temperature τ. 0: Top-K expert indices K∈1,…,ET×KT_K∈\1,…,E\^T× K, where row K(xt)=TopK((xt))T_K(x_t)=Top\!K( s(x_t)) over the proposal set tC_t is the selected set for token t in Eq. 7; normalized gate probabilities π~∈ℝ≥0T×K π _≥ 0^T× K with ∑i∈K(xt)π~t,i=1 _i _K(x_t) π_t,i=1; per-expert token counts n∈ℕEn ^E. Throughout, i∈1,…,Ei∈\1,…,E\ identifies an expert and t∈1,…,Tt∈\1,…,T\ a token. 1: W~r←RowNorm(Wr,cr) W_r← RowNorm(W_r,c_r) ℓ2 _2-normalize each router row to target norm crc_r 2: Θ←HW~r⊤ ← H\, W_r^\! router logits 3: s′←σ(Θ)s ←σ( ) clean scores, no bias 4: s←s′+ϕs← s +φ LFLB bias-added selection scores, Eq. 18 in Appendix A.3 5: for each token t∈Tt∈ T (parallel) do 6: gt,i∼Gumbel(0,1)g_t,i (0,1) for i=1,…,Ei=1,…,E i.i.d. standard Gumbel per expert 7: gt←gt−1T∑t′=1Tgt′g_t← g_t- 1T\! _t =1^Tg_t per-expert noise centering 8: t←Top−Mi(logst,i′+τgt,i)C_t← *Top-\!M_i ( s _t,i+τ\,g_t,i ) proposal on clean scores 9: K,t←Top−Ki∈t(st,i)T_K,t← *Top-\!K_i _t (s_t,i ) bias-aware correction 10: ζt←∑j∈K,tst,j′ _t← _j _K,ts _t,j 11: π~t,i←st,i′ζt,i∈K,t π_t,i← s _t,i _t, i _K,t gates use clean scores, ∑=1Σ=1 12: end for 13: ni←∑t=1T[i∈K,t],i=1,…,En_i← _t=1^T1\! [i _K,t ], i=1,…,E entries of the per-expert token-count vector n 14: return K,π~,nT_K,\ π,\ n 3.3 Routing 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 of wang2024lossfree, 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 affinity si′(x)=σ(θi(x))s _i(x)=σ( _i(x)) with i denoting an expert and θi(x)=w~r,i⊤x _i(x)= w_r,i^\! x denoting the router logit. It selects the top-K experts by a bias-adjusted score si(x)=si′(x)+ϕis_i(x)=s _i(x)+ _i, and combines their outputs with the clean scores si′(x)s _i(x) renormalized over the selected set (wang2024lossfree; deepseekai2024v3). The per-expert bias ϕi _i is a controller state nudged once per training step toward under-used experts from the observed per-expert token counts nin_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 in Appendix A.3) 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 and detailed in Appendix A.3. The algorithm takes as input a batch of T tokens with their pre-MoE hidden states H, the router weights WrW_r, and the current per-expert bias ϕφ, and returns the top-K expert indices KT_K selected per token, their normalized gate probabilities π~ π, and the per-expert token counts n that 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-K selection involves a Gumbel proposal-correction sampler. It entails a stochastic top-M proposal on the clean sigmoid scores followed by a bias-aware top-K correction (M>KM>K, with M a pre-selected hyperparameter fixed across all runs), with the per-expert noise centered across the batch. The top-M proposal followed by the top-K correction separates exploration from final routing. We use Gumbel scores only to expand the candidate set to M, then choose the final K using clean scores inside that set. Since applying Gumbel noise over all experts can cause large arbitrary flips, we restrict it to a top-M candidate pool among E experts. It lets near-boundary experts get exposure while preventing far-tail experts from winning due to noise. Lines 5–9 of Algorithm 1 follow this procedure. • The router rows are ℓ2 _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 of deepseekai2024v3, described in Appendix A.3. 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. 3.4 Sparsity S and Expert Split Factor G In considering the geometry of MoE models, the two architectural dimensions of interest are the total expert count E and the per-token activated experts K. Together with LlayersL_layers, d, and dffd_f, these determine an active parameter budget NactN_act (the model FLOPs-relevant parameter count seen by any individual token) and a total parameter budget NtotN_tot (the memory footprint of the model). In this section, we discuss two architectural knobs, the sparsity S and the expert split factor G, that we later use to parametrize the sparse MoE scaling laws. We recall from Table 1 that the two parameter budgets set the sparsity S=1−Nact/NtotS=1-N_act/N_tot, the fraction of parameters left inactive per token. To compare architectures whose experts differ in size, we use expert split factor defined as: G≔dffdexpertG\;\; \;\; d_fd_expert (8) It is the ratio of the dense FFN width dffd_f that an MoE block replaces to a single expert’s hidden width dexpertd_expert, measuring how finely that dense FFN has been partitioned into experts. We visually show the construction of this expert split factor G in an MoE layer in Figure 14 in Appendix B. dffd_f is the intermediate (hidden) width of the reference full FFN, a fixed multiple of the residual-stream width dmodeld_model, and is distinct from dmodeld_model itself and dexpert=dffGd_expert= d_fG is the per-expert slice of it. The G=1G\!=\!1 case recovers the setting studied in clark2022unified in which each expert is itself a full-width FFN. G>1G\!>\!1 corresponds 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 higher G means smaller, but more finer-grained experts in the MoE layer. The ratio EG EG is the effective expansion rate of total to dense FFN parameters used by krajewski2024scaling, so increasing G at fixed EG EG yields more, smaller experts at constant total parameter count. 4 Sparse MoE Scaling Laws ℒ(.)L(.) Building on the sparse MoE architecture discussed in Section 3, we elaborate on what Z entails and we lay out the desiderata behind the scaling law formulations we target in our study. We want a loss law where Z meets 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 the u free dimensions in a scaling law is evaluated at v discrete values using a full factorial design, the number of configurations scales as O(vu)O(v^u). The laws fitted in the study by abnar2025parameters and by tian2026towards are 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 sparsity S and expert split factor G described in Section 3, and argue that G is a necessary axis beyond those considered in abnar2025parameters. We then review the abnar2025parameters law on our training runs in Section 4.1, following which we discuss the final scaling function form we use for the optimization in Section 4.2. (a) G=4G=4. (b) G=8G=8. Figure 3: Predicted-vs-observed training loss for the ℒ(Ntot,S,D)L(N_tot,S,D) fit in Equation (9), fit per expert split factor. The panels for every stratum we fit, G∈1,2,4,8G\!∈\!\1,2,4,8\, are in Figure 15 of Appendix B. 4.1 Scaling Laws with Sparsity Before fitting the functional form following the desiderata above, we follow the functional form studied in abnar2025parameters, 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_tot,S), ℒ(Ntot,S,D)=aNtotα⏟capacity+bDβ⏟data+c(1−S)λ⏟sparsity floor+j(1−S)δNtotγ⏟interaction+eL(N_tot,S,D)\;=\; aN_tot^α_capacity\;+\; bD^β_data\;+\; c(1-S)^λ_sparsity floor\;+\; j(1-S)^δ\,N_tot^γ_interaction\;+\;e (9) with sparsity S, which differs from how sparsity (=1−KE=1- KE) is defined in that study. The first two terms recover the Chinchilla form (hoffmann2022training) in the limit S→0S\!→\!0 (dense model). We fit the functional form in Equation (9) separately for each expert split factor G, 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 between 0.00020.0002 and 0.00370.0037 in loss units across the four expert split factors, under half a percent relative error. Figure 3 shows G=4G\!=\!4 and G=8G\!=\!8, and the full G∈1,2,4,8G\!∈\!\1,2,4,8\ panel is given in Appendix B.1. What the parity plots also reveal is that the fitted coefficients change substantially across G, so no single set of (Ntot,S,D)(N_tot,S,D) coefficients describes all the expert split factors at once. This is what leads us to build G into the law in Section 4.2. In Figure 4, we find that three trends are monotone over S∈[0.5,0.981]S∈[0.5,0.981]. These three trends agree with those abnar2025parameters report, 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. Figure 4: Model FLOPs optimal Ntot⋆N_tot , Nact⋆N_act , and ℒ⋆L as a function of sparsity S, at five budgets C. The curves are computed from our own G=4G\!=\!4 fit of Equation (9) on our runs. Larger MoE sparsity S increases optimal total parameters, decreases optimal active parameters, and lowers achievable loss at every model-FLOPs budget. 4.2 G-augmented Joint Scaling Law Before adding the expert split factor G to the law, we note that G and the sparsity S are distinct knobs that happen to share the parameter counts. The two move different geometry variables and appear coupled only because both these dimensions modulate NtotN_tot and NactN_act. Empirically they are close to orthogonal on our grid, where the Pearson correlation between logG G and each of logNtot N_tot, S, and logNact N_act is at most 0.020.02 in magnitude, because the sweep varied dexpertd_expert independently of (E,K)(E,K) and every G∈1,2,4,8G\!∈\!\1,2,4,8\ is observed across nearly the full NtotN_tot and S ranges. A starting point to understand how to augment the scaling functional form with the expert split factor G is the conventional model FLOPs optimal recipe. Figure 5 shows that the loss curves in NtotN_tot fan out rather than stay parallel as G or S varies, so the effect of the expert split factor and of sparsity is not a rigid shift but depends on NtotN_tot, which is a coupling a separable law in (Ntot,S,D)(N_tot,S,D) cannot express. Following this observation, we tie G, S, and NtotN_tot together in a single interaction term. The joint scaling law we use, on the scaling law dimensions (Ntot,S,G)(N_tot,S,G) induced by Z, is ℒ(Ntot,S,D,G) (N_tot,S,D,G) =aNtotα+bDβ+c(1−S)λ \;=\; aN_tot^α\;+\; bD^β\;+\; c(1-S)^λ +j(1−S)δNtotγGη+e. +\; j(1-S)^δ\,N_tot^γ\,G^η\;+\;e. (10) Here G appears only in the denominator of the interaction term, so raising G shrinks the sparsity penalty and lowers the predicted loss at fixed (Ntot,S,D)(N_tot,S,D). Placing G in the joint coupling rather than folding it into a single effective parameter count is what keeps the fit balanced across G values. Figure 5: Loss against total MoE parameters NtotN_tot. In the left, one curve per expert split factor G at fixed S=0.95S\!=\!0.95, D=6D\!=\!6 B. In the right, one curve per sparsity band at fixed G=4G\!=\!4, D=6D\!=\!6 B. The curves fan out rather than staying parallel (left), which is why G enters the law as its own axis instead of being absorbed into an effective parameter count, and Table 2 bears this out, where the G-exponent η is identified rather than sitting on the coefficient ridge. Coeff. Value 95%95\% BCa CI Identifiability a 8.4628.462 [8.24, 8.85][8.24,\,8.85] ridge b 88.3288.32 [87.8, 89.9][87.8,\,89.9] ridge c 0.18280.1828 [0.181, 0.184][0.181,\,0.184] ridge j 0.71230.7123 [0.710, 0.713][0.710,\,0.713] ridge e 0.31290.3129 [0.307, 0.316][0.307,\,0.316] ridge α 0.10480.1048 [0.104, 0.106][0.104,\,0.106] ridge β 0.20700.2070 [0.206, 0.208][0.206,\,0.208] identified λ 0.12490.1249 [0.119, 0.138][0.119,\,0.138] ridge δ 0.55670.5567 [0.550, 0.577][0.550,\,0.577] ridge γ 0.17020.1702 [0.167, 0.174][0.167,\,0.174] ridge η +0.9513+0.9513 [+0.911,+0.954][+0.911,\,+0.954] identified Table 2: Fitted coefficients of the joint scaling law in Equation (10). 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. Fitting Results. We fit the form in Equation (10) on the scaling grid spanning G∈1,2,4,6,8G\!∈\!\1,2,4,6,8\. We hold out the top-10%10\% of runs by NactN_act within each G-stratum to validate the observations we do not use in training. We minimize the Huber objective on logℒ (δHuber=0.1 _Huber=0.1) by L-BFGS-B in log-coefficient space. Table 2 reports the point estimates together with 95%95\% confidence intervals, and marks each coefficient as identified or ridge. Figure 6: Training loss against expert split factor G=dff/dexpertG\!=\!d_f/d_expert at matched (Ntot,S,D)(N_tot,S,D) operating points. The predicted loss decreases monotonically with G. The observed runs decrease up to G=4G\!=\!4 and flatten or slightly reverse at G=8G\!=\!8 within 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 shows that η and β are identified here, because G and D each enter the law in a single place that no other coefficient can absorb. The G exponent η≈0.95η\!≈\!0.95 is therefore our robust observation, and since it multiplies the joint sparsity and capacity term, its reading is that larger G discounts 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 space G∈1,2,4,6,8G\!∈\!\1,2,4,6,8\. We check this against the runs in Figure 6, where the predictions of Equation (10) decrease monotonically with G at matched (Ntot,S,D)(N_tot,S,D) while the observed losses decrease up to G=4G\!=\!4 and flatten or slightly reverse at G=8G\!=\!8 within 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)^-δ factor rather than to NtotN_tot itself, since at matched S the Ntot−γN_tot^-γ decay of the interaction term shrinks the gap as models grow. From the systems viewpoint of this paper, η quantifies 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 will price against sparse MoE layer communications overhead. We further validate the fit quality in Figure 7, which collapses a subset of all runs across G∈1,2,4,6,8G\!∈\!\1,2,4,6,8\ onto the identity line, and the joint law attains a held-out RMSE of 0.007790.00779 in loss units on the largest 10%10\% of runs by NactN_act within each G stratum. The fitted form also behaves as we expect of a well-formed scaling law. It is monotone-decreasing in NtotN_tot, D, and G, and convex in S with a sharp penalty as S→1S\!→\!1. As G→∞G\!→\!∞ the joint term drops out and the loss settles onto a/Ntotα+b/Dβ+c(1−S)−λ+ea/N_tot^α+b/D^β+c\,(1-S)^-λ+e. As Ntot,D→∞N_tot,D\!→\!∞ at fixed MoE sparsity S, the loss approaches e+c(1−S)−λe+c\,(1-S)^-λ, so the sparsity floor persists at infinite scale and the irreducible term is e∞=e+c≈0.50e_∞\!=\!e+c\!≈\!0.50 in the dense limit S→0S\!→\!0 rather than e alone. Figure 7: Joint-law parity plot over all runs, R2=0.99877R^2=0.99877. The crosses are the held-out validation runs, the largest 10%10\% by NactN_act within each G stratum, with a held-out RMSE of 0.007790.00779 in loss units. Marker shape and color pick out a sample of the (G,K)(G,K) pairs the sweep covers. 5 Estimating MFU(.)MFU(.) with a Performance Model This section specifies the design of a performance estimator that predicts MFU(Z,Pℓ,B)MFU(Z,P_ ,B) and Mem(Z,Pℓ,B)Mem(Z,P_ ,B) for MOSAIC in Section 2. 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.1 Design 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.2 Functionality The performance model is driven by a model-architecture specification that fully describes a training run, as shown in Appendix E. It groups settings into seven blocks. The model block 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. The model.moe block adds the mixture-of-experts layer (expert count, top-k routing, expert width, shared-expert width, and the capacity factor that sizes the routed token budget and the all-to-all traffic). The search block declares the device budget. The performance block sets the activation-checkpointing strategy (none, selective, super-selective, or full). The optimizer block selects the optimizer family (adam, adamw, muon, scion, and more), which sets the optimizer-state footprint and update FLOPs. The data block gives the batch and sequence shape, and the hardware block names the target node type. From this specification, the tool builds an internal model and hardware representation, then enumerates the feasible execution plans Pℓ=P_ = (PTPP_TP, PEPP_EP, PCPP_CP, P_P, PDPP_DP, BmicroB_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, where we also describe the layout search that produces the ranking. 5.3 Microbenchmarks 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 fused GroupedLinear operator, keyed by the number of local experts (num_gemms), the per-expert token load (the grouped tile M), 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 tile Mexpert=⌈TseqBmicroK/E⌉M_expert= T_seq\,B_micro\,K/E (sequence length TseqT_seq, micro-batch BmicroB_micro, top-K, E experts), 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-k routing and element-wise operations (silu, residual_add), forward and backward. • Cross-entropy over the vocabulary: the fused vocab-parallel loss on fp32-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, and send_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 to 8×8× in 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.4 Iteration-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.1 ms 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=nmbtmb,nmb=BglobalPDPBmicrot_mb= \! (t_cmp,\ t_disp ), t_block=n_mb\,t_mb, n_mb= B_globalP_DP\,B_micro (11) where nmbn_mb is the number of microbatches per optimizer step, tcmpt_cmp is the summed measured GPU-compute time of the stage’s layers, and tdisp=Lstageκ(Eloc)tgap+thostt_disp=L_stage\,κ(E_loc)\,t_gap+t_host is the CPU-dispatch floor over LstageL_stage layers per stage, with a measured per-kernel gap tgapt_gap, a fixed per-microbatch host term thostt_host, and a kernel-count law κ(Eloc)=κ0+κ1Elocκ(E_loc)= _0+ _1\,E_loc that grows with the number of local experts ElocE_loc (the surrounding framework kernels, namely permutation, copies, and bias-adds, dominate the launch count, not the fused expert GEMM). The max makes 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. For nmbn_mb microbatches per pipeline, P_P pipeline stages, and PVPPP_VPP-way interleaving 444https://docs.nvidia.com/nemo/megatron-bridge/0.2.0/parallelisms.html#interleaved-pipeline-parallel-schedule, the bubble fraction is fbubble=P−1PVPPnmbf_bubble= P_P-1P_VPP\,n_mb (12) so the pipeline bubble time is tbubble=fbubbletblockt_bubble=f_bubble\,t_block. Each rank computes nmbn_mb microbatches of work and idles (P−1)/PVPP(P_P-1)/P_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 the 2(P−1)2(P_P-1) warmup/cooldown microbatches a (P−1)/P(P_P-1)/P_P stage bubble on top of the (P−1)/nmb(P_P-1)/n_mb steady-state term, and clamped the result at (P−1)/P(P_P-1)/P_P. 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 surplus tlastt_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_vocab=(n_mb-1)\, (t_first,t_last)+t_first+t_last (13) with an added tlastnmbfbubblet_last\,n_mb\,f_bubble term for the ramp. Here tfirstt_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, with the DP grad reduce-scatter and param all-gather overlapped against compute where windows allow. 5.5 Calibration 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=χsystiteranalytict^cal_iter= _sys\,t^analytic_iter (14) where the systems calibration multiplier is χsys=1+cA2AtA2Atiter 1[NnodeEP>1]+cPP(NnodePP−1)+covl 1[P≥Ngpu/node,NnodePP=1]. split _sys&=1+c_A2A\, t_A2At_iter\,1\! [N^EP_node>1 ]+c_P\, (N^P_node-1 )\\ & +c_ovl\,1\! [P_P≥ N_gpu/node,\ N^P_node=1 ]. split (15) Here NnodeEP=⌈PEPPTP/Ngpu/node⌉N^EP_node= P_EP\,P_TP/N_gpu/node and NnodePP=⌈P/Ngpu/node⌉N^P_node= P_P/N_gpu/node count how many nodes of Ngpu/nodeN_gpu/node GPUs each the expert-parallel and pipeline-parallel groups span, and tA2At_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. (i) The deep-pipeline term scales with the number of inter-node hops the pipeline crosses. (i) The intra-node overlap credit (covl<0c_ovl<0) applies only to a deep pipeline confined to one node. Single-node, shallow-pipeline configurations incur no correction (χsys=1 _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 the MFUMFU and MemMem predictors that MOSAIC consumes. We validate these predictions against measured MFU on real training sweeps in Section 7.2. 6 Optimal MoE Sparsity Scaling In Section 2, we defined deliverable model FLOPs CdeliverableC_deliverable as the systems-aware constraint used by MOSAIC. We now develop the motivation sketched in Figure 1 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 FLOPs CdeliverableC_deliverable of Equation (2) 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=6NactDC_model\!=\!6\,N_actD. Simply said, we want the (Ntot,S,D,G)(N_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_tot ,S ,D ,G ), treating these as continuous variables. In the second stage, we map that optimum to a discrete, deployable model configuration, since S and G are not free knobs one can set directly. S is a derived ratio, and G moves only by changing the discrete geometry (E,K,Llayers,d)(E,K,L_layers,d) of Section 3. To that end, Sections 6.1 and 6.2 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.1 Model-FLOPs Optimal Sparsity Is Boundary-Seeking Conditioning on the scaling law, the conventional scaling studies seek the optimal sparsity S that minimizes loss at a fixed budget CmodelC_model. Concretely, with Nact=(1−S)NtotN_act=(1-S)\,N_tot and the budget Cmodel=6NactDC_model=6\,N_act\,D (at fixed G), the model FLOPs only optimization problem reads minNtot,S,Dℒ(Ntot,S,D,G)s.t.6(1−S)NtotD=Cmodel,Ntot,D>0,S∈[0,1). _N_tot,S,D\;L(N_tot,S,D,G) .t. 6\,(1-S)\,N_tot\,D\;=\;C_model, N_tot,D>0, S∈[0,1). (16) Solving Equation (16) with the joint law of Equation (10) at Cmodel∈3×1019, 6×1019, 1020, 3×1020, 6×1020C_model∈\3× 10^19,\,6× 10^19,\,10^20,\,3× 10^20,\,6× 10^20\ FLOPs recovers the three trends that the G=4G\!=\!4 fit of Equation (9) produces in Figure 4. The important observation in the plots is that Ntot⋆(S)N_tot (S) increases monotonically with S, Nact⋆(S)N_act (S) decreases, and ℒ⋆(S)L (S) decreases monotonically at every CmodelC_model. Because ℒ⋆(S)L (S) decreases monotonically over the calibrated sparsity range, the interior condition ∂ℒ⋆/∂S=0 /∂ S=0 is not satisfied within this region. The optimal sparsity under the CmodelC_model budget, therefore lies at S=SmaxS=S_ , the upper boundary of our data support. Thus, within this range where the scaling law is calibrated, model-FLOPs only optimization reduces to “make S as large as the supported architecture space allows.” The left column of Figure 8 shows the same monotone descent on both knobs, in S at fixed G and in G at fixed S. We do not claim that this monotonic trend continues arbitrarily close to S=1S=1, where the asymptotic sparsity-penalty terms in the fitted law eventually dominate. Figure 8: Opposing views of model scaling and systems cost. The top row varies MoE sparsity S holding G=4G\!=\!4. The bottom row varies the expert split factor G holding S=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 parameters NtotN_tot under that budget. The right column shows the wall-clock time each configuration needs per delivered model FLOP on 512512 devices (64 p6 nodes), relative to the training time of the lowest sparsity configuration. The CmodelC_model optimal loss configuration pays 1.70×1.70× per FLOP in the top row and 2.66×2.66× in the bottom. The starred winners sit at opposite ends, so the CmodelC_model optimal loss configuration is the systems-slowest. 6.2 The Boundary Optimum Is Not Necessarily the Cluster Optimum In the preceding section, the model-FLOPs optimum proved boundary-seeking in sparsity, returning the largest S the 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. The CdeliverableC_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. 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 provide CdeliverableC_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 shows how loss, sparsity, and execution cost trade off on the way to the deployable optimum. Its left panel plots the optimal loss ℒ⋆L reached at each sparsity level at fixed Cmodel=1022C_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-size 512512 (number of GPU devices) with the performance model of Section 5, optimized over the parallelism layout, and the wall-clock per model FLOP rises monotonically with sparsity to 1.70×1.70× the densest configuration. This opposing ordering is precisely why the architecture-dependent deliverable model-FLOPs constraint introduced in Section 2 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 of abnar2025parameters that at a fixed total parameter count a larger training budget lowers the optimal sparsity, we hold neither NactN_act nor NtotN_tot fixed in Figure 8. We instead propose a shift toward hardware constraints of Section 2 to resolve these optimality questions. 7 Experiments 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 in Appendix D 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. Sparsity S and the expert split factor G are 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)L(N_tot,S,D,G) of Equation (10), not the per-G fits of Equation (9). After presenting the ladder, we validate the performance model against measured MFU in Section 7.2, 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 fixed CpeakC_ peak envelope in Section 7.3, (2) the optimal MoE-layer geometry in sparsity and expert split factor G as we vary the cluster size (nodes) and training window (days) in Section 7.4, and (3) staged real training runs that test the predicted rank ordering against measured throughput in Section 7.5. 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 with 88 NVIDIA B200 GPUs. We provide the details of the Megatron and Torch-based framework we stage our runs on in Appendix F.3, 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.1 Geometry 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 follow Llayers=mqL_layers=m\,q and nhead=mqn_head=m\,q with the multiplier m fixed to 4, so the residual width is d=nheaddheadd=n_head\,d_head at dhead=128d_head=128, and the reference dense FFN width is dff=2.5d_f=2.5\,d. A single seed q thus fixes depth (LlayersL_layers), width (d), 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 sparsity S, are searched on top of it. We fix the co-scaling functions of our ladder this way and vary the geometry seed q, and we leave other variants of co-scaling geometry as future work. With the ladder in place, the framework searches over the seed q, the expert count E, the top-K, the split factor G, and the data budget D, returning the optimum, and the sparsity S and the parameter counts NtotN_tot and NactN_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 of Appendix F. 7.2 Performance Model Validation Our validation spans three sweeps on these p6 nodes, over the q=4q\!=\!4, q=8q\!=\!8, and q=12q\!=\!12 rungs of Table 9. The q=4q\!=\!4 model (0.70.7B NactN_act, 1616B NtotN_tot) runs on 4 nodes (3232 GPUs), the q=8q\!=\!8 model (5.45.4B NactN_act, 127127B NtotN_tot) on 32 nodes (256256 GPUs) and 64 nodes (512512 GPUs), and the q=12q\!=\!12 model (1818B NactN_act, 434434B NtotN_tot) on 64 nodes (512512 GPUs). Each data point pairs one parallelism configuration, namely expert (ep), pipeline (p), 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. We vary the global batch size and restrict tp=1tp=1 and cp=1cp=1 for efficiency. We release the estimator as ScalePlan666https://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 feasP_feas of Equation (C-FLOPs), so the predictions validated here and every estimator number reported later can be regenerated from a single model and cluster specification. Figure 9: Predicted versus Real MFU per parallelism configuration, each panel normalized to its own maximum and sorted by real MFU. 7.2.1 MFU 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 under 15%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 is 100%100\% for every 5.4B and 18B sweep and 80%80\% for the 700M sweep. Top-10 overlap is 100%100\% for all 5.4B and 18B sweeps and 90%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.2 Implications 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. The dominant lever is not any single degree but the pipeline fill ratio rfill=nmb/PPPr_fill=n_mb/P_P, which explains almost all of the up to 3.1×3.1× layout spread in measured MFU, with throughput collapsing as rfillr_fill falls toward 11 and saturating for rfill≥4r_fill≥ 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 fix PTP=1P_TP=1, set PEPP_EP to the scale-up domain, and choose the (P,Bmicro)(P_P,B_micro) pair that keeps rfill≥4r_fill≥ 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.3 Optimal Allocation Under CpeakC_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. Inputs Hardware 32 p6 nodes (Ndev=256N_dev=256 devices) Per-device peak FpeakF_peak 2.252.25 PFLOPs (bf16/dense) Wall-clock TtrainT_train 20 days Per-device memory cap MdevM_dev 179179 GB Goodput ηgood _good 1.01.0 Raw hardware ceiling CpeakC_peak 9.95×10239.95× 10^23 FLOPs Global / sequence batch Bglobal=1024B_global=1024, Tseq=2048T_seq=2048 Loss law joint ℒ(Ntot,S,D,G)L(N_tot,S,D,G) from Eq. (10) Search space (output) Outer architecture grid q,E,K,Gq,E,K,G Inner layout search over PTPP_TP, PCPP_CP, PEP,P,PDPP_EP,P_P,P_DP, Bmicro∈1,2,4B_micro\!∈\!\1,2,4\, and activation-checkpointing ∈none, selective, super-selective, full∈\!\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 of Z, and the inner parallelism search realizes the embedded maximization over feasP_feas in Equation (C-FLOPs). We target a deployment simulation with 32 p6 nodes (Ndev=256N_dev=256 devices) and Ttrain=20daysT_train=20\,days. These inputs are listed in Table 3. The model FLOPs the returned optimum realizes under this ceiling lie about 150×150× beyond our largest fit run, an extrapolation the deployed law supports for the configuration ordering. Combining the scaling ladder of Table 9 with the sparse MoE axes, we restrict the free discrete search dimensions to q q ∈8,10,12,14,16,18, ∈\8,10,12,14,16,18\, K K ∈1,2,4,6,8, ∈\1,2,4,6,8\, E E ∈4,6,8,16,24,32,48,96,128,192,256,384,512, ∈\4,6,8,16,24,32,48,96,128,192,256,384,512\, G G ∈1,2,4,8. ∈\1,2,4,8\. We add a practical geometry guard to prune the grid by setting S≥0.7S≥ 0.7, since the scaling law data we collect to fit the loss law lies predominantly in the S≥0.7S≥ 0.7 region, 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. We fix the token budget D per configuration by sweeping the tokens-per-parameter ratio TPP=D/NactTPP=D/N_act over the grid TPP∈[40,150]TPP∈[40,150]. Our training data spans TPP up to 7070 (Table 6), so the cap sits about twice beyond the calibrated range, and we keep only (Z,TPP)(Z,TPP) pairs whose CmodelC_model fits the deliverable model FLOPs CdeliverableC_deliverable at that layout’s MFU together with the memory and parallelism constraints. 7.3.1 Results MOSAIC returns a single loss optimal configuration, reported in Table 4. The constrained optimum is a q=12,Llayers=48,E=96,K=2,G=4q\!=\!12,L_layers\!=\!48,E\!=\!96,K\!=\!2,G\!=\!4 model with Nact=14.5N_act\!=\!14.5 B at sparsity S=0.956S\!=\!0.956, running at 12.4%12.4\% predicted MFU and loss ℒ=1.3898L\!=\!1.3898. Its selected layout is P=8P_P=8, PEP=8P_EP=8, PCP=1P_CP=1, PTP=1P_TP=1 with micro-batch Bmicro=4B_micro=4. LlayersL_layers E K G NactN_act (B) S TPP D (T) MFU (%) Loss p ep mb cp tp 48 96 2 4 14.5 0.956 98 1.42 12.38 1.3898 8 8 4 1 1 Table 4: The loss optimal configuration MOSAIC returns under the envelope of Table 3. The MFU column is the predicted peak MFU of the best feasible execution plan, not a measured value. Although Table 4 reports only the loss optimal configuration, we now show all feasible (Z,D)(Z,D) configurations under the same envelope in Figure 10. While D is 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 of 9.95×10239.95× 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. It plots the loss against the model FLOPs CmodelC_model each candidate realizes, color-coded by G, with the gold star marking the loss optimal configuration. Figure 10: Loss versus model FLOPs CmodelC_model over all feasible (Z,D)(Z,D) configurations under the 3232-node, 2020-day envelope, color-coded by G. Every candidate draws on the same raw hardware ceiling of 9.95×10239.95× 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 only by sustaining a higher MFU, which turns the fixed 2020 days of wall-clock into more useful model FLOPs. The expert split factor G orders that spread. The G=8G\!=\!8 geometries sit at the low CmodelC_model end at a median MFU near 8%8\%, and the coarser G=2G\!=\!2 geometries reach the right end near 1313–20%20\%, so a higher G delivers 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). 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. The candidate that converts the shared hardware ceiling CpeakC_peak into the most model FLOPs is a denser G=2,E=32G\!=\!2,E\!=\!32 geometry at 1.94×10231.94× 10^23 FLOPs, and it still lands at a higher loss than the optimal configuration. The loss optimal configuration is a G=4,E=96,K=2G\!=\!4,E\!=\!96,K\!=\!2 model at Nact=14.5N_act\!=\!14.5B and S=0.956S\!=\!0.956 that reaches ℒ=1.3898L\!=\!1.3898 on 1.23×10231.23× 10^23 model FLOPs, roughly 36%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.4 Optimal S and G 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 at 3232 p6 nodes (256256 GPUs) and varies the wall-clock window over 10,20,30,40\10,20,30,40\ days. The second fixes the window at 3030 days 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. (a) 3232 p6 nodes. (b) 3030 days. Figure 11: Optimal predicted loss vs. the budget under the two sweeps. Markers are colored by the optimal expert split factor G⋆G and annotated with S⋆S and Ntot⋆N_tot . Dashed lines denote best per G configuration. The solid line is the overall optimum. The optimal sparsity S⋆S does not move monotonically with the budget in either sweep, staying on a wide plateau of roughly 0.9150.915 to 0.9630.963. We plot the results from running MOSAIC in Figure 11. We observe that G⋆G is 44 across most budgets, with G=8G\!=\!8 appearing only at the longest window and G=2G\!=\!2 at the largest node budget, and the exact G⋆G marker is often a near-tie. Looking at the span of sparsity, we find that S⋆S sits on a wide plateau (roughly 0.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. Finally, Nact⋆N_act moves slowly with budget, staying near 1414–1515 B parameters across the days range and growing to 3434 B and 5858 B 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 high S and high G. 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.5 From 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 (6464 GPUs) for 1010 days, 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 lists the selected configurations with the MFU the framework predicts for each. rank (loss) LlayersL_layers E K G NactN_act (B) NtotN_tot (B) S pred. MFU (%) norm. actual MFU p ep cp tp mbs ID 1 40 128 2 4 8.5 250.0 0.966 14.0 0.84 4 8 1 1 4 ID 2 40 128 2 8 5.6 126.4 0.956 11.6 0.69 4 8 1 1 4 ID 3 40 32 2 4 8.5 66.0 0.872 21.8 0.91 4 8 1 1 4 ID 4 32 32 2 2 7.3 65.8 0.889 23.3 1.0 4 8 1 1 4 Table 5: The four configurations on the 88-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) in Appendix A.2. 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. 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 the absolute accuracy on which the feasibility constraint and the token budget rest. Normalizing the predicted MFU the same way gives 0.600.60, 0.500.50, 0.940.94, and 1.001.00 against the measured 0.840.84, 0.690.69, 0.910.91, and 1.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 divides by the measured MFU of each run. The largest model FLOPs these curves reach is roughly 1515–20×20× the largest budget in the scaling law fitting set of Appendix C, which tops out at 8.1×10208.1× 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, 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.0 against 65.865.8 B) and close sparsity (0.8720.872 against 0.8890.889), yet the G=4G\!=\!4 run stays below the G=2G\!=\!2 run on both axes. On the per-FLOP axis, up to roughly 102210^22 model FLOPs the G=8G\!=\!8, low-NactN_act configuration (ID 2) holds an advantage over ID 1, since its smaller active count (5.65.6 B against 77–8.58.5 B) 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. Figure 12: Smoothed loss trajectories of the staged runs against model FLOPs CmodelC_model (left) and against the peak-equivalent hardware compute consumed Cmodel/MFUC_model/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 (at 102210^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 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_cost=C_model/(MFU\, _good), the compute proportional to the GPU-hours a run spends to deliver CmodelC_model, and we hold ηgood=1 _good=1 throughout, so CcostC_cost reduces to Cmodel/MFUC_model/MFU as plotted. Dividing out MFU charges each run for the hardware time it actually spends, and ID 1 sustains the higher MFU (normalized 0.840.84 against ID 2’s 0.690.69), so the same hardware buys ID 1 about 1.2×1.2× the model computation. Figure 12 shows this directly, where ID 2 leads on the model-FLOPs panel (left) while ID 1 leads on the CcostC_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, 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. 8 Related Work MoE Scaling Laws. Past work on sparse MoE architecture extends Chinchilla-style scaling laws (hoffmann2022training) to sparse MoEs. krajewski2024scaling introduce granularity as an explicit axis and abnar2025parameters characterize the sparsity dimension, while ludziejewski2025joint fit a joint law over active size, experts, and tokens and argue MoEs can be more memory-efficient than dense models. wan2026holistic push toward holistic recipes that map a budget to a full architecture, and mcleish2025gemstones show 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, fernandez2024hardware measure 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. bian2025inference target inference-efficient model selection, and on the MoE side huang2023towards and zheng2026uniep attack 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. anthony2024codesign align model dimensions to GEMM and wave-quantization boundaries whereas bian2026architecture jointly tune architecture for accuracy and inference efficiency. bae2025hybrid studies hybrid architectures for compute optimality that are better on recall and wang2025catransformers co-design transformers and accelerators under a carbon objective. We also see studies like zhang2025ladder which 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 as lescao2022million ablate architecture at fixed GPU-hours, and wu2024inference add 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. 9 Conclusion, 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 factor G may itself be an artifact of that stack’s systems penalty. Thus, broader kernel-level changes could turn the optimal G into 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 to G≤8G≤ 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 and dffd_f co-scaling families would test whether the (Ntot,S,G)(N_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 A Pre-training Recipe In this section, we cover some more details of the pretraining recipe and the architecture that went into the training runs. A.1 Architecture 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 of wang2024lossfree in the router. In our instantiation, the shared expert has the same hidden width dexpertd_expert as a routed expert, so as the expert split factor G increases, dexpertd_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 router WrW_r, and the embedding and readout matrices, are drawn flat from (0,0.0062)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/2Llayers)2)N\! (0,(0.006/ 2L_layers)^2 ), so deeper models initialize the residual-write weights smaller and keep the residual-stream variance stable across the 2Llayers2L_layers residual additions per forward pass. For the normalization layers, their gains are set to 11 and biases to 0. Initializing WrW_r like any other fan-in matrix leaves router logits approximately uniform across experts at step zero. We leave μ -style initialization for hyperparameter transfer (jiang2026hpmoe) in sparse MoE models as future work. (a) Loss vs. peak learning rate per configuration, with the per-configuration optimum (stars) read from a quadratic fit in log10lr _10lr. (b) Optimal learning rate vs. model FLOPs Cmodel=6NactDC_model=6\,N_act\,D, solid line is the fit lr⋆(Cmodel)=3.30Cmodel−0.199lr (C_model)=3.30\,C_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 FLOPs CmodelC_model, with one law predicting the optimal learning rate across the full range of total parameters and sparsity. A.2 Optimizer and Learning-rate Settings We use AdamW at weight decay 0.10.1 with the warmup-stable-decay (WSD) schedule of hagele2024scaling, namely a fixed warmup of approximately 100100-million tokens regardless of run length, a stable phase at peak learning rate, and a linear decay to zero over the final 20%20\% training steps. The AdamW moments are the standard defaults β1=0.9 _1=0.9 and β2=0.999 _2=0.999 with ε=10−14 =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 FLOPs Cmodel=6NactDC_model=6\,N_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 in log10lr _10lr (Figure 13, 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). Fitting log10lr⋆ _10lr against log10Cmodel _10C_model over these near-optimal points gives lr⋆(Cmodel)= 3.30Cmodel−0.199,Cmodel=6NactDlr (C_model)\;=\;3.30\;C_model^-0.199, C_model=6\,N_act\,D\, (17) with a tight parallel band around the central line (Figure 13(b)) 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 on CmodelC_model, we detect no systematic residual dependence of the near-optimal learning rate on NtotN_tot, sparsity, or the Nact/DN_act/D split. 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, where tian2026towards report a slope of −0.153-0.153 over a comparable compute range, close to our value. In practice, this gives a single rule, where for any candidate configuration we compute Cmodel=6NactDC_model=6\,N_act\,D and read the peak learning rate off Equation (17). Out-of-scope HPO axes. Optimizer hyper-parameters (weight decay 0.10.1, warmup, scheduler shape), sequence length (Tseq=2048T_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.3 MoE Routing Details In this section, we describe the choices behind the routing algorithm for our scaling recipe presented in Algorithm 1. Sigmoid scoring with per-expert bias. Following wang2024lossfree, we instantiate the score vector (x) s(x) of Equation (7) as a clean sigmoid score si′(x)s _i(x) plus an additive per-expert bias ϕ∈ℝEφ ^E, so that si(x)=si′(x)+ϕi=σ(θi(x))+ϕi,θi(x)=w~r,i⊤xs_i(x)\;=\;s _i(x)+ _i\;=\;σ\! ( _i(x) )+ _i, _i(x)\;=\; w_r,i^\! \,x (18) where w~r,i w_r,i denotes the row-normalized i-th row of WrW_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 router z-loss or logit capping. The bias ϕφ is 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, let nin_i denote the number of tokens routed to expert i in a batch and let n¯=1E∑i=1Eni n= 1E _i=1^En_i. Following wang2024lossfree; team2025kimi, we update ϕφ in 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+μbuiu_i\;=\; n-n_i E^-1 _j( n-n_j)^2+ , _i\;←\; _i+ _b\,u_i (19) We additionally center the bias after each update, ϕi←ϕi−1E∑jϕj _i← _i- 1E _j _j, since only relative bias differences affect routing. Probabilistic top-K via Gumbel proposal–correction. We replace deterministic selection with a Gumbel-noise proposal–correction sampler operating at a fixed temperature τ : (Gumbel proposal)t (Gumbel proposal) _t =TopMi(logst,i′+τgt,i), =Top\!M_i ( s _t,i+τ\,g_t,i ), gt,i g_t,i ∼Gumbel(0,1), (0,1), (20) (Gumbel correction)K,t (Gumbel correction) _K,t =TopKi∈t(st,i), =Top\!K_i _t (s_t,i ), M M =min(ρK,E). = (ρ K,E). (21) The construction rests on the Gumbel-max theorem. For K=1K\!=\!1 and gi∼Gumbel(0,1)g_i (0,1), argmaxi(logpi+τgi)∼Categorical(pi1/τ∑jpj1/τ) *arg\,max_i ( p_i+τ\,g_i )\; \;Categorical\! ( p_i^1/τ _jp_j^1/τ ) (22) and its top-K extension gives exact sampling without replacement from the same distribution (yellott1977relationship; kool2019stochastic). Injecting the noise in log-score space, as in Equation (20), therefore draws the proposal set exactly from P(i)∝st,i′1/τP(i) s 1/τ_t,i on the clean scores, and the bias ϕφ is then applied in the correction stage of Equation (21) 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 makes the gt,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 is M=min(ρK,E)M= (ρ K,E) with K≤EK≤ E, since ρKρ K exceeds E at the small expert counts in our grid. The temperature τ sets how peaked the proposal distribution is, and as τ→0τ\!→\!0 the algorithm approaches a deterministic clean-score top-M proposal followed by the bias-aware top-K correction within that proposal. This limit need not coincide with the global deterministic top-K on st,is_t,i, since the bias can favor an expert outside the clean top-M. For the uncentered, uncorrected Gumbel-max baseline with two experts and K=1K\!=\!1, the probability of selecting b over a is σ(τ−1log(sb′/sa′))σ\! (τ^-1 (s _b/s _a) ), so near-tied experts receive meaningful exploration while confident decisions are exponentially protected. In Algorithm 1 itself, at two experts ρ=12ρ=12 puts both in the proposal and the correction then selects deterministically under st,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 g_t,i=g_t,i- 1T _tg_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τ=0.01 and ρ=12ρ=12, with token-level granularity across all our runs, and we use μb=10−3 _b=10^-3 and ε=10−4 =10^-4 in the bias update of Equation (19). We find that this method of sampling with low τ allows 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 parameters NactN_act and more. Compared to adding more dense layers at the beginning as done in singh2026arcee which reduces the model capacity at matched NactN_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 of WrW_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 pass w~r,i=cr⋅wr,i‖wr,i‖2 w_r,i\;=\;c_r· w_r,i\|w_r,i\|_2, where the target norm crc_r is a fixed constant, in our case set to 1.0. It allows to recover routing decisions that depend on directional alignment between wr,iw_r,i and x rather than on accidental norm drift. This is the RowNorm(Wr,cr) RowNorm(W_r,c_r) step in Algorithm 1, which produces the normalized router W~r W_r used to compute the pre-sigmoid logits Θ . Additionally, the bias update of Eq. (19) is invoked once per optimizer step on the per-expert counts n returned 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 of shazeer2017outrageously; 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, we apply the router z-loss of zoph2022st to suppress unbounded growth in the pre-sigmoid router logits Θt,i=θi(xt) _t,i= _i(x_t) of Eq. (18): ℒz=1T∑t=1T(log∑i=1EeΘt,i)2L_z\;=\; 1T _t=1^T ( _i=1^Ee _t,i )^2 (23) added to the cross-entropy loss with a small fixed coefficient. Because selection in our routing variant is taken on σ(Θ)+ϕσ( )+φ rather than on Θ directly, ℒzL_z acts purely as a numerical guard on the logits and does not interfere with the bias controller. We set the coefficient of z-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 in deepseekai2024v3. We set its coefficient to 1e-3 in all our runs. Appendix B Scaling Analysis Figure 14: Example of an MoE layer in a transformer block, at three expert split factors G∈1,4,8G\!∈\!\1,4,8\. In every panel, the routed experts sit beside a shared expert. dffd_f is the hidden width of a single reference dense feed-forward expert (G=1G\!=\!1), and the total routed model width EdffE\,d_f is held fixed across the panels. Raising G splits each reference FFN into G narrower experts, so the routed-expert count grows (E→GE\!→\!G\,E) while each expert’s hidden width shrinks proportionally (dexpert=dff/Gd_expert=d_f/G). Relative to G=1G\!=\!1, a factor-G configuration has G times as many experts, each G times narrower. The top-K selection is kept the same across the three panels, so the active parameters NactN_act shrink with higher G rather than staying fixed. Note that this is different from granularity defined in krajewski2024scaling. This appendix collects the per-G evidence behind the scaling laws of Section 4, starting with Figure 14, which shows how the expert split factor G of Section 3 reshapes an MoE layer. (a) G=1G=1. (b) G=2G=2. (c) G=4G=4. (d) G=8G=8. Figure 15: Predicted-vs-observed loss for the ℒ(Ntot,S,D)L(N_tot,S,D) fit with Equation (9), refit independently on each expert split factor stratum G∈1,2,4,8G\!∈\!\1,2,4,8\. The G=4G\!=\!4 and G=8G\!=\!8 panels are the ones reproduced in Figure 3. B.1 Per-G Scaling Law Fits This appendix section gives the per-G details behind Section 4.1. We fit ℒ(Ntot,S,D)L(N_tot,S,D) in Equation (9) on each G separately, holding out the largest 10%10\% by active parameters. The fit is noise-limited in all the per G optimization of the functional form and the held-out large models match it, with no residual structure on E, K, NtotN_tot, S, D, or LlayersL_layers (all Pearson correlations below 0.130.13 in magnitude). Refitting the same form independently per expert split factor as shown in Figure 15, (G∈1,2,4,8G\!∈\!\1,2,4,8\) fits each G well in isolation, but the fitted coefficients drift substantially across G, so no single set of (Ntot,S,D)(N_tot,S,D) coefficients describes all expert split factors at once. This is the signal that G must enter the law itself, which we do in Section 4.2. B.2 Joint-Law Uncertainty and Identifiability This appendix section gives the evidence behind the G-augmented joint law of Section 4.2, namely that a larger expert split factor G lowers 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η\!≈\!+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 η, we perform 500500 warm-started wild-bootstrap refits using Rademacher residual perturbations and report a 95%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 of 400400 fresh 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 η and β are 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 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,α), the joint-term parameters (j,γ,δ)(j,γ,δ), and the stand-alone sparsity floor (c,λ)(c,λ) trade off against one another at near-constant objective. This is why we report only η and β as identified and treat the rest as ridge coordinates: the law’s loss value and its per-G calibration are robust, but no single split of the prefactors and exponents is privileged, and pinning one for publication requires fixing a prefactor (for example a) and refitting. The stand-alone sparsity-floor term c/(1−S)λc/(1-S)^λ is also numerically small at our S≤0.981S\!≤\!0.981 because most of the sparsity penalty is carried by the joint j/((1−S)δNtotγGη)j/((1-S)^δ\,N_tot^γ\,G^η) term. Finally, the asymptotic blow-up ℒ→∞L\!→\!∞ as S→1S\!→\!1 is imposed by the form rather than measured. Pinning (c,λ,δ)(c,λ,δ) empirically would require runs at S near 0.9990.999, which our grid does not contain. Appendix C Scaling Sweep Training Runs Quantity Min Median Max Total parameters NtotN_tot (B) 0.22 5.34 79.24 Active parameters NactN_act (B) 0.104 0.339 2.72 Sparsity S 0.499 0.944 0.981 Expert split factor G 1 - 8 Expert count E 8 - 512 Top-K activation 1 - 8 LlayersL_layers 8 - 24 Residual width d 1024 - 3072 Training tokens D (B) 2.88 5.76 56.38 TPP=D/NactTPP=D/N_act 10.8 - 70 Compute (non-embedding FLOPs) 1.9×10181.9× 10^18 - 8.1×10208.1× 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, trained for over 200B tokens each, are not included. G NtotN_tot (B) NactN_act (B) S D (B) LlayersL_layers E, K 1 0.54–30.91 0.152–1.179 0.72–0.98 3.84–28.19 8,12,16 E∈8–128E\!∈\!\8 128\, K∈1,2K\!∈\!\1,2\ 2 0.95–12.65 0.124–0.387 0.87–0.97 3.84–5.76 8,12,16 E∈32–128E\!∈\!\32 128\, K∈1,2K\!∈\!\1,2\ 4 0.22–79.24 0.110–2.723 0.50–0.98 2.88–56.38 8,12,16,24 E∈12–384E\!∈\!\12 384\, K∈2,4,6,8K\!∈\!\2,4,6,8\ 8 0.71–53.0 0.104–2.52 0.85–0.98 3.84–50.00 8,12,16,24 E∈96–512E\!∈\!\96 512\, K∈2,4,8K\!∈\!\2,4,8\ Table 7: Ranges of total and active parameters, sparsity, training tokens, depth, expert count, and top-K spanned by the scaling runs at each expert split factor G, where B denotes billions. C.1 Coverage of the Scaling-Law Dataset This appendix documents the slice of the (Ntot,Nact,S,G,D)(N_tot,N_act,S,G,D) space that the joint law of Section 4.2 is fit on. The dataset comprises nruns=150n_runs=150 from-scratch sparse MoE pretraining runs. Our runs span Ntot∈[0.22,79.24]N_tot∈[0.22,79.24]\,B parameters, Nact∈[0.10,2.72]N_act∈[0.10,2.72]\,B parameters, S∈[0.499,0.981]S∈[0.499,0.981], G∈1,2,4,6,8G∈\1,2,4,6,8\, E∈8,12,16,24,32,48,64,96,128,192,256,384,512E∈\8,12,16,24,32,48,64,96,128,192,256,384,512\, K∈1,2,4,6,8K∈\1,2,4,6,8\, training tokens D∈[2.88,56.38]D∈[2.88,56.38]\,B (TPP=D/Nact∈[10.8,70]TPP=D/N_act∈[10.8,70]), across four architecture families (Llayers,d)∈(8,1024),(12,1536),(16,2048),(24,3072)(L_layers,d)∈\(8,1024),(12,1536),(16,2048),(24,3072)\. Total compute for our scaling law fits spans 1.9×10181.9× 10^18–8.1×10208.1× 10^20 FLOPs. We display some of the statistics of our training runs in Table 6. LlayersL_layers d nrunsn_runs D (B) NtotN_tot (B) G values 8 1024 52 2.88–7.68 0.22–7.14 1,2,4,6,8\1,2,4,6,8\ 12 1536 56 5.76–11.52 0.73–25.20 1,2,4,6,8\1,2,4,6,8\ 16 2048 39 28.19–56.38 4.13–45.66 1,2,4,8\1,2,4,8\ 24 3072 4 50.00 40.16–79.24 4,8\4,8\ Table 8: Per-family coverage. D is chained to (Llayers,d)(L_layers,d) by an outer Chinchilla-style budgeting decision. C.2 Training Dynamics of the MoE Models Figure 16 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 than 5%5\% of the tokens in a batch that a uniform assignment over the E experts 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 keeps routing stable and balanced throughout training even at the largest scales. Figure 16: Training dynamics of the four largest MoE models in our scaling run sweep (annotated with NtotN_tot, NactN_act, E, K, G). 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 D MOSAIC Optimization D.1 Sparse-MoE Execution-Plan Realization For the Megatron-Core stack used in our experiments, the dense layers factor the world size as Ndev=PTPPCPPPDPN_dev=P_TP\,P_CP\,P_P\,P_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=PEPPEDP,PEDP=PDPPEPP_DP\;=\;P_EP\,P_EDP, P_EDP\;=\; P_DPP_EP (24) which requires that PEPP_EP divides PDPP_DP. Equivalently, the MoE-layer mesh factors as Ndev=PTPPCPPPEPPEDPN_dev=P_TP\,P_CP\,P_P\,P_EP\,P_EDP. Thus PEPP_EP is an input degree of the plan that participates in the MoE-layer device factorization, but it does not multiply the dense-layer PDPP_DP axis independently, and the expert data parallelism PEDPP_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)P_disc(N_dev) of Section 2.2. D.2 Structural Properties Three properties of Equation (5) are directly useful for interpreting the result. Property 1 (objective convexity). In the continuous variables (logNtot,logD,log(1−S),logG)∈ℝ4( N_tot, D, (1-S), G) ^4, the loss ℒ(Ntot,S,D,G)L(N_tot,S,D,G) of Equation (10) is convex, since ℒL is a posynomial in (Ntot,D,1−S,G)(N_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) resides in the constraints, not in the objective. Property 3 (decoupling of D from layout). Tokens D enter only the objective and the FLOPs budget, never the memory cap or the layout product. Consequently, for any Z and any Pℓ∈feas(Z,B)P_ _feas(Z,B), the largest admissible D is the closed-form D⋆(Z,Pℓ)=FpeakNdevTtrainMFU(Z,Pℓ,B)ηgood(Ndev)6Nact(Ntot,S)D (Z,P_ )\;=\; F_peak\,N_dev\,T_train\,MFU(Z,P_ ,B)\, _good(N_dev)6\,N_act(N_tot,S) (25) and since ℒL is monotone-decreasing in D, any optimum saturates Equation (C-FLOPs) at D=D⋆(Z,Pℓ)D=D (Z,P_ ). This eliminates one continuous decision variable analytically and reduces the inner problem to a search over (Z,Pℓ)(Z,P_ ) alone. Algorithmic consequence. Properties 1–3 are what let the bilevel program of Equation (5) be solved exactly by the single-stage search of Section 7, 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_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 in D, so an optimal run never leaves wall-clock budget unused and Equation (25) fixes the saturating D⋆D in closed form, removing the only continuous decision variable. The deployed search of Section 7 realizes this on the finite grid TPP∈[40,150]TPP∈[40,150], taking for each candidate the largest ratio the budget supports, so D⋆D is reached up to grid resolution and the cap binds only for the smallest models. What remains splits cleanly, because at fixed Z the best execution plan is simply the feasible plan of maximum MFU within Equation (4), 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, fixing D from Equation (25), and scoring the loss with Equation (10), and the configuration it returns is the exact optimum of Equation (5) over the realizable grid rather than a heuristic solution. A generic MINLP solver would in any case have little to work with here, since MFUMFU and MemMem come out of the calibrated performance model of Section 5 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.3 Optimal 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 of hoffmann2022training would look like, and what assumptions would be required to state one. Nothing below replaces the discrete search of Section 7. D.3.1 Reduction to a CpeakC_peak Law The exogenous hardware-time budget is the raw ceiling Cpeak=FpeakNdevTtrainC_peak=F_peakN_devT_train of Section 2.1 together with the goodput fraction ηgood(Ndev) _good(N_dev). Both are fixed before an architecture is selected, so CpeakηgoodC_peak _good is the counterpart of the compute budget in a Chinchilla-style prescription, whereas the deliverable budget CdeliverableC_deliverable is not, since it also carries the best attainable MFU⋆(Z)=maxPℓ∈feas(Z,B)MFU(Z,Pℓ,B)MFU (Z)= _P_ _feas(Z,B)MFU(Z,P_ ,B) over the feasible plan set feas(Z,B)P_feas(Z,B) of Section 2.2. Loss decreases monotonically in the token count, so saturating Equation (C-FLOPs) recovers Property 3, D⋆(Z,Cpeak)=CpeakMFU⋆(Z)ηgood6Nact(Z).D (Z,C_peak)= C_peak\,MFU (Z)\, _good6N_act(Z). (26) Since MFU⋆(Z)MFU (Z) does not depend on the training-window length, substituting Equation (26) into Equation (10) collapses the law to ℒ⋆(Z,Cpeak)=ℒ∞(Z)+Ψ(Z)Cpeakβ,Ψ(Z)=b(6Nact(Z)MFU⋆(Z)ηgood)βL (Z,C_peak)=L_∞(Z)+ (Z)C_peak^β, (Z)=b ( 6N_act(Z)MFU (Z)\, _good )^\!β (27) where ℒ∞(Z)L_∞(Z) collects the three terms of Equation (10) that do not involve D, together with the irreducible constant e. We note that D⋆∝CpeakD C_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 in CpeakC_peak because the selected architecture can change the constant MFU⋆(Z)/Nact(Z)MFU (Z)/N_act(Z). D.3.2 Conditions 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 write r≡1−Sr≡ 1-S and N≡NtotN≡ N_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)MFU (N,r,G) is differentiable. We define the local MFU log-slopes ωN=∂logMFU⋆/∂logN _N=∂ /∂ N, ωS=∂logMFU⋆/∂logr _S=∂ /∂ r, and ωG=−∂logMFU⋆/∂logG _G=-\,∂ /∂ G, so that ωG>0 _G>0 means MFU falls as experts become finer. After writing the four terms of Equation (10) as ℒcap=aN−αL_cap=aN^-α, ℒdata=bD−βL_data=bD^-β, ℒsp=cr−λL_sp=cr^-λ and ℒint=jN−γr−δG−ηL_int=jN^-γr^-δG^-η, and eliminating D through Equation (26), stationarity in (logN,logr,logG)( N, r, G), we obtain αℒcap α\,L_cap =βℒdata(1−ωN−γωG/η), =β\,L_data (1- _N-γ _G/η ), λℒsp λ\,L_sp =βℒdata(1−ωS−δωG/η),ηℒint=βωGℒdata. =β\,L_data (1- _S-δ _G/η ), η\,L_int=β\, _G\,L_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, the G equation requires ωG>0 _G>0, while the N and r equations require 1−ωN−γωG/η>01- _N-γ _G/η>0 and 1−ωS−δωG/η>01- _S-δ _G/η>0, respectively. Without a systems penalty on finer experts, ωG=0 _G=0, the third stationarity condition cannot hold. The objective is then strictly decreasing in G within this regime, so no interior G optimum 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) 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 that MFU⋆≃m0(N/N0)ωN(r/r0)ωS(G/G0)−ωGMFU m_0(N/N_0) _N(r/r_0) _S(G/G_0)^- _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⋆)(α−γ−αδ/λ)/η \;r (N )^α/λ, D (N )^α/β, G (N )^(α-γ-αδ/λ)/η\; (29) Writing p=α/λp=α/λ, ϱ=α/β =α/β and g=(α−γ−αδ/λ)/ηg=(α-γ-αδ/λ)/η for these exponents, and ν=[(1−ωN)+(1−ωS)p+ϱ+ωGg]−1ν=[(1- _N)+(1- _S)p+ + _Gg]^-1, the same path in terms of the purchased budget is N⋆∝Cpeakν,1−S⋆∝Cpeakpν,D⋆∝Cpeakϱν,G⋆∝Cpeakgν. \;N C_peak^ν, 1-S C_peak^pν, D C_peak ν, G C_peak^gν.\; (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 E Performance Model Input/Outputs The listings below describe ScalePlan (https://github.com/dmlc/ScalePlan), the performance model released with this paper. Beyond scoring one plan, it enumerates the execution plans Pℓ=(PTP,PEP,PCP,P,PDP,Bmicro,Ackpt)P_ =(P_TP,P_EP,P_CP,P_P,P_DP,B_micro,A_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 feasP_feas of Equation (C-FLOPs), 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: 2 n_layers: 48 # L_layers, transformer blocks 3 hidden_sz: 6144 # d, residual-stream width 4 inter_sz: 15360 # d_f, dense FFN width (2.5 d) 5 n_q_heads: 48 # n_head, query heads 6 n_kv_heads: 8 # key/value heads (grouped-query attention) 7 head_dim: 128 # d_head, so d = n_head * d_head 8 vocab_sz: 257536 # vocabulary size 9 precision: "bf16" 10 sdpa_precision: "bf16" 11 glu: true 12 rotary_embeds: true 13 dropout: false 14 tie_embeddings: false 15 16 moe: 17 n_experts: 128 # E, routed experts per MoE block 18 experts_per_token: 3 # K, top-K active experts 19 capacity_factor: 1.0 20 expert_inter_sz: 3840 # d_expert, so G = d_f / d_expert = 4 21 moe_frequency: 1 22 expert_tp_degree: 1 # P_ETP, expert tensor parallel degree 23 24search: 25 num_devices: 512 # N_dev 26 27performance: 28 activation_checkpointing_type: full 29 30optimizer: 31 optimizer_type: muon 32 33data: 34 gbs: 8192 # B_global, sequences per optimizer step 35 seqlen: 8192 # T_seq, sequence length 36 microbatch_sz: 2 # B_micro 37 38hardware: 39 node_type: p6-b200.48xlarge # 8 GPUs per node, so N_gpu/node = 8 Listing 1: Performance model input specification. Output example, where the layout degrees are (PTP,PEP,P,PCP,PDP)(P_TP,P_EP,P_P,P_CP,P_DP) of PℓP_ and the reported MFU and memory are the MFU(Z,Pℓ,B)MFU(Z,P_ ,B) and Mem(Z,Pℓ,B)Mem(Z,P_ ,B) predictions of Section 5. ⬇ 1Rank 1: 2 Configuration: tp=1, ep=8, p=16, cp=1, dp=32 3 Analytical MFU: 13.69% 4 Memory per device: 21.30 GB 5 Iteration time: 44.892 s Listing 2: Performance model output. Appendix F Results F.1 Geometry Scale Ladders Table 9 lists the rungs of the geometry scaling ladder of Section 7.1 that our search enumerates. Each row is one seed q, giving the backbone dimensions the seed fixes and, at a representative MoE configuration, the parameter counts and sparsity that backbone realizes. q LlayersL_layers nheadn_head d dffd_f NtotN_tot (B)∗ NactN_act (B)∗ S 4 16 16 2048 5120 15.5 0.7 0.954 6 24 24 3072 7680 53.2 2.3 0.957 8 32 32 4096 10240 127.3 5.4 0.958 10 40 40 5120 12800 250.0 10.4 0.958 12 48 48 6144 15360 433.7 17.9 0.959 14 56 56 7168 17920 690.5 28.2 0.959 16 64 64 8192 20480 1032.8 41.9 0.959 18 72 72 9216 23040 1472.8 59.5 0.960 Table 9: Seeds of the model-scale ladder used in our experiments. The geometry seed q sets the backbone depth Llayers=4qL_layers=4q, head count nhead=4qn_head=4q, residual width d=128nheadd=128\,n_head, and reference dense FFN width dff=2.5d_f=2.5\,d. The last three columns give the total and active parameter counts and the induced sparsity S at an example MoE configuration E=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.2 Implications for Sparse-MoE Parallelism Design This appendix develops the design rules summarized in Section 7.2.2. 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 ratio rfill=nmb/PPPr_fill=n_mb/P_P, the microbatches in flight per pipeline stage. Within a sweep, layout choice moves measured MFU by up to 3.1×3.1× (18B) and 2.1×2.1× (700M), and almost all of that spread tracks rfillr_fill, since 1F1B leaves each stage idle for a bubble fraction of (P−1)/(PVPPnmb)(P_P-1)/(P_VPP\,n_mb), so throughput collapses as rfillr_fill falls toward 11 and saturates once rfillr_fill is large. The 18B pipeline sweep makes this explicit, where at P=16P_P=16, MFU climbs as rfillr_fill goes 0.25→0.5→10.25→ 0.5→ 1, then plateaus at ∼10 \!10–12%12\% for rfill≥4r_fill≥ 4. This reframes the two parallelism knobs that people tune independently, since deeper pipelines and larger microbatches both shrink rfillr_fill and are only free once the fill ratio is already comfortable. Concretely, every extra pipeline stage or microbatch that drops rfillr_fill below ∼4 \!4 costs measurable MFU, whereas the same change at rfill≥4r_fill≥ 4 is 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 with PEPP_EP in the 700M sweep (for PEP∈1,2,4,8P_EP∈\1,2,4,8\ at fixed depth) and drops from PEP=8P_EP=8 to PEP=16P_EP=16 in both 5.4B sweeps as the group leaves the 8-GPU node. Once PEPP_EP is pinned to the domain, moderate pipeline depth is the cheapest remaining axis provided it does not push rfillr_fill below the plateau, and the 5.4B optima use P∈2,4P_P∈\2,4\ because at these depths rfillr_fill stays ≥8≥ 8. This is the actionable ordering the model exposes, where we set PEPP_EP to the scale-up domain, then add pipeline depth only while rfillr_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 fix PTP=1P_TP=1, set PEPP_EP to the scale-up domain, then choose the (P,Bmicro)(P_P,B_micro) pair that keeps the pipeline fill ratio rfill≥4r_fill≥ 4 at 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.3 The Training Stack Behind the Staged Runs The staged runs of Section 7.5 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/moe as 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). 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 full 20482048-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 to 14.414.4 TB/s of bidirectional intra-node NVLink and a fourth-generation Elastic Fabric Adapter delivering 3.23.2 Tbps (400400 Gbps 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.