Frontier Pre-training at Microsoft AI (MAI-Base-2)

ListenThe Recipe Behind a 2.8-Trillion-Parameter Base Model10:56

Show notes

Microsoft AI's lab notes on MAI-Base-2 — a 2.8-trillion-total, 80-billion-active MoE base model trained on 16K GB200s in 38 days — read like a frontier-scale cookbook. We unpack the three co-designed workstreams: scaling science (a joint learning-rate-times-weight-decay law), model–infrastructure co-design (dropless multi-round MoE, distributed Muon), and data scaling (encoder-free multimodal pretraining). Along the way we pull out what any image or video modeling team can steal, and where the authors are careful about what they actually show. **Sources** - Frontier Pre-training at Microsoft AI (MAI-Base-2) — The Microsoft AI Team

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

MAI-Base-2 is Microsoft AI’s 2.8T-total / 80B-active MoE base model, pre-trained on 16K GB200 GPUs over 38 days (21% MFU, 89% goodput) and reported to beat every tested open-weight base on held-out BPB evaluations while using 4.2–12.3× fewer pre-training FLOPs than DeepSeek-V4-Pro-Base. The post is a methodology write-up of three co-designed workstreams — scaling science (source-scaled ladders, joint LR×WD×dropout laws), model-infrastructure co-design (dropless multi-round MoE with tunable fixed memory, resharding-based distributed Muon, activation gradient clipping), and data scaling (36% more organic tokens than MAI-Base-1, native encoder-free vision pretraining). Framed as the successor to the 1T/35B MAI-Thinking-1 (MAI-Thinking-1: Building a Hill-Climbing Machine) and the base for MAI-Base-3 and the next post-trained reasoning model. The explicit stance — organic data only, no distillation from other models — is a frontier-scale continuation of MAI’s “capabilities should be learned, not inherited” position.

Key claims

2.8T total / 80B active parameter MoE with “MoE in every layer” (vs MAI-Thinking-1’s interleaved dense/MoE), trained on 50T tokens including 4.5T unique VLM tokens; natively multimodal via an encoder-free VLM architecture [§“How MAI-Base-2 Compares”, §“Scaling Data”].
Reported 4.2–12.3× FLOPs-efficiency over DeepSeek-V4-Pro-Base and 6.3–12.7× over MAI-Base-1 on four held-out evaluations (Code, STEM, Math, Q&A), estimated by fitting an independent BPB-vs-FLOPs scaling law per evaluation and reading off the FLOPs needed to match the reference [§“Isolating modeling efficiency from compute”, Fig. 3].
Scaling ladders use source scaling: on lower rungs, the set of unique samples is artificially restricted so each dataset is repeated the same number of times regardless of training horizon, keeping unique-data-to-parameters ratio fixed across the ladder. Without it, small-scale experiments favor under-regularized configurations that overfit at scale [§“Robust scaling ladders”].
Joint LR × WD scaling law fit on a (log LR, log(LR·WD)) grid rather than (log LR, log WD) because LR·WD is the sensitivity axis: LR* · WD* = 17.04 / steps and log LR* = -1.51 + 0.175 log TPP - 0.573 log steps. The ~-1 steps exponent for LR·WD matches the view that the decoupled weight-decay time scale should be a fixed fraction of training; the ~-0.5 steps exponent for LR matches the √2-batch-doubling rule [§“Scaling laws for hyperparameters”].
Dropout fitted as sigmoid(0.526 (log ActiveParams - 27.2)) — i.e. a scaling law for dropout as a function of active-parameter count [§“Scaling laws for hyperparameters”].
Co-design metric: efficiency gain EG = baseline FLOPs or wall-clock to reach candidate's eval loss, reported in both EGFLOPs and EGtime variants. EGtime guides shipping decisions; EGFLOPs wins that lose on wall-clock mark where to invest in infrastructure [§“Codesigning our infrastructure and MAI-Base-2”].
MoE-every-layer was EGFLOPs-optimal but EGtime-suboptimal in MAI-Thinking-1 due to dynamic-routing memory fragmentation and per-expert overhead. For MAI-Base-2, fixing both lifted MoE-layer MFU from 24% → 35%, flipping the EGtime verdict toward MoE-every-layer [§“Scaling MoE”].
Multi-round dropless MoE: experts are grouped per EP rank; each group runs one or more dispatch → compute → collect rounds with a fixed per-expert token capacity. Overflow tokens use unused capacity from other experts in the same group; otherwise an extra round. Memory usage is static regardless of token imbalance; capacity is tunable, with steady-state configured for a single round. Backward pass uses fine-grained per-expert-group recompute; dispatch/compute/collect pipelined so all but the first dispatch and last collect overlap with expert compute [§“Scaling MoE”, Fig. 6].
Per-expert overhead cut via: grouped GEMM + NVLink per expert group, device-resident token metadata + GPU-initiated all-to-all, symmetric-memory all-to-all over NVLink without staging copies, and batched load-balancing-loss all-reduce across layers [§“Reducing per-expert overhead”].
Distributed Muon reshards momentum during the optimizer step: FP32→BF16 shard update, parameter-shape-batched all-to-all into full matrices, symmetry-aware Newton–Schulz locally, second all-to-all back to shards. Overlaps via multi-batch prefetching, caches resharding metadata, and replicates (not shards) MoE router weights to skip the first all-to-all. Final cost: 2.1% / 165 ms per training iteration vs theoretical 55 ms communication-bound peak [§“Scaling Muon”].
Activation gradient clipping replaces global weight-gradient-norm clipping: clips the gradients of activations at selected network locations, so outlier gradients from a small portion of a batch don’t scale down well-behaved examples. Keeps gradient norm in a much tighter range than MAI-Base-1, and its benefit grows with scale — unusual, since most training changes yield a roughly constant efficiency gain across scales [§“Activation gradient clipping”, Figs. 7, 8].
Activation gradient clipping doubles as a data-validation diagnostic: traced a clipped example back to a malformed code pull-request where a license header was cut off mid-sentence, revealing a preprocessing bug [§“Activation gradient clipping”, Fig. 9].
Pre-training corpus grew ~36% in unique training tokens vs MAI-Base-1; expansions include quality classification, topic binning, web parsing, deduplication, evaluation decontamination, and embedding-based diversity sampling to reduce duplicative content. March 2026 knowledge cutoff [§“Scaling Data”, Fig. 10].
VLM corpus is 4.5T unique tokens (18% of the 50T training mix with repetition), covering image captioning, chart understanding, visual grounding, spatial reasoning, medical understanding, computer use, sheet-music reading [§“Scaling Data”].
Three stated design principles: intelligence should be learned from first principles, simple methods scale best, scientific rigor over shortcuts [§“What’s next”].

Method

MAI-Base-2’s development is organized around a scaling-ladder-driven ablation protocol where each rung is a model at a fixed aspect ratio parameterized by layer count L (hidden size = (256/3)·L, query heads = L rounded to nearest 16). Candidate recipe changes are evaluated by fitting BPB-vs-FLOPs and BPB-vs-time power laws on the ladder and extracting the FLOPs/time ratio at matched loss (EG metric). Source scaling ensures each dataset is repeated the same number of times at every ladder rung, so smaller-scale ablations transfer upward without rewarding under-regularized configurations that would overfit at full scale.

Three concrete architectural/infrastructure contributions underpin the headline FLOPs-efficiency: (1) the dropless multi-round MoE that keeps memory static under arbitrary token routing imbalance while allowing MoE-in-every-layer to be EGtime-optimal rather than only EGFLOPs-optimal; (2) a distributed Muon implementation that reshards momentum per-step via two batched all-to-alls, enabling full-matrix Newton–Schulz orthogonalization under ZeRO sharding at 16K-GPU scale with 2.1% iteration overhead; (3) activation gradient clipping as a scale-growing stability lever, replacing global weight-gradient norm clipping with per-activation local clips at selected network locations.

The hyperparameter recipe is derived from a joint (log LR, log(LR·WD)) grid fit at multiple (model size, TPP) combinations with 2D-quadratic surfaces, yielding closed-form LR*(TPP, steps) and WD*(steps) laws plus a sigmoid dropout law in active parameters. Together with the vision-side encoder-free VLM pretraining on 4.5T unique multimodal tokens (18% of 50T), these constitute the complete recipe used for the final 38-day training run on 16K GB200s.

Results

2.8T total / 80B active MoE outperforms every tested open-weight model (DeepSeek-V4-Pro-Base, Kimi K3, Qwen3.8-2.4T-A95B, MiMo-V2.6-Pro-RL) on held-out BPB across Code, STEM, Math, Q&A [Fig. 1].
4.2–12.3× FLOPs-efficiency vs DeepSeek-V4-Pro-Base; 6.3–12.7× FLOPs-efficiency vs MAI-Base-1 [Fig. 3].
MoE MFU lifted from 24% to 35% via dropless multi-round MoE + per-expert overhead reductions, flipping EGtime verdict toward MoE-in-every-layer [§“Scaling MoE”].
Distributed Muon: 2.1% (165 ms) of average iteration time at 16K GPUs vs 55 ms theoretical communication-bound peak [§“Scaling Muon”].
Training run: 16K GB200 GPUs, 38 days (August–September 2026), 21% MFU (normalized by 2.5 PFLOPS BF16, excluding rematerialization and optimizer), 89% goodput [§“How MAI-Base-2 Compares”].
Pre-training mix: 50T tokens total; 4.5T unique VLM tokens; ~36% more unique training tokens than MAI-Base-1 [§“Scaling Data”].

Why it’s interesting

This is the first filed frontier-scale case study where joint LR × WD scaling as a single LR·WD-steps law replaces the usual decoupled search, giving explicit closed-form prescriptions (LR·WD = 17.04/steps, log LR = -1.51 + 0.175 log TPP - 0.573 log steps) that extend the Fitting paradigm this wiki tracks under Hyperparameter scaling laws — complementary to How to Set the Learning Rate for Large-Scale Pre-training?‘s $\eta^*(N, D)$ surface and to Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization‘s Hyperball weight-norm sphere, both of which also attack the hyperparameter-search-is-the-dominant-cost problem but from the single-axis or geometric-constraint angle rather than the joint-coordinate-system angle. The multi-round dropless MoE with tunable fixed per-expert memory is a cleanly specified instance of the “non-uniformity as first-class primitive” thesis in Distributed training parallelism, and the shift from MAI-Thinking-1’s interleaved dense/MoE to MoE-in-every-layer — justified via EGtime after lifting MoE MFU from 24% to 35% — gives a sharp counter-datapoint to the FLOPs-vs-wall-clock trade-off story in MAI-Thinking-1: Building a Hill-Climbing Machine (same project, different conclusion once the systems bottleneck was removed). Activation gradient clipping as a scale-growing efficiency lever is unusual on this page’s board (most changes give roughly constant EG), and its double use as a data-validation diagnostic — tracing a clipped example back to a preprocessing bug — generalizes the “optimizer-signal-as-diagnostic” pattern in Mitigate Silent Expert Death In Ultra-Sparse MoE and Data Scarcity and Model Sparsity: Mixtures-of-Experts Overfit More to Repeated Data.