The Training and RL Infrastructure Behind MAI-Thinking-1

MAI-Thinking-1 is Microsoft AI’s first in-house reasoning model, a mixture-of-experts (MoE) model with 35B active and 1T total parameters, post-trained from the base model MAI-Base-1. This note collects what the MAI-Thinking-1 technical report says about the systems side: the GPU fleet, the numerical recipe, the YOLO training framework, the Rocket reinforcement learning (RL) stack, reliability, and scheduling. It leaves out the data, the RL recipe and the evaluations.

Everything here comes from the report. I have not run any of it. Where I add arithmetic or an inference of my own, I say so.

1. Hardware and cluster layout#

MAI-Thinking-1 was developed on a mixed fleet of NVIDIA H100, GB200 and GB300 systems. The GB200 and GB300 clusters sit in Microsoft’s own datacenters. Production serving runs on Microsoft’s MAIA-200 accelerator.

Phase	Hardware	Site	Tokens and context
Pre-training	8,192 GB200	Phoenix, AZ	30T tokens at 16K
Mid-training 1	8,192 GB200	Phoenix, AZ	3.4T tokens at 64K
Mid-training 2	4,096 GB200	Phoenix, AZ	150B tokens at 256K
RL climb	about 4,600 GB300	Dallas, TX	rollouts up to 128K

Pre-training ran on a single logical cluster at one site. The report gives two sizes for RL: the climb used “4.6K GB300s”, and the largest RL job used 4,864 GB300s.

Racks. The GB200 and GB300 systems are NVL72 racks, each a 72-GPU NVLink domain. Racks talk to each other over InfiniBand with GPUDirect RDMA.
NVL64. From model version 2 onward, training jobs used only 64 GPUs per rack and kept the rest as spare capacity for failed nodes and unhealthy devices. That leaves 8 GPUs per rack. My inference: a GB200 NVL72 rack holds 18 compute trays of 4 GPUs, so 64 GPUs is 16 trays with 2 spare. That matches the certification rule in section 6, which excluded racks with fewer than 16 healthy nodes.
H100. The 8-GPU H100 nodes (NVLink/NVSwitch inside a node, InfiniBand between nodes) stayed in the lab environment, alongside GB200 and GB300 systems, for development, validation, comparative profiling and next-generation bring-up. Model versions v0 and v1 were early prototypes trained on H100.
MAIA-200. Against a GB200 deployment, MAIA-200 serves MAI-Thinking-1 with over 40% higher token generation throughput under the same rack power budget.

The model the systems run#

Several infrastructure choices below follow from the shape of MAI-Base-1:

Property	Value
Parameters	34.7B active, 962B total (rounded to 35B / 1T in the report)
Layers, hidden size	78, 6,656
Attention	5 sliding-window layers (window 512, RoPE) per global layer (no position encoding); GQA with 8 KV heads and 80 query heads, head dimension 128; RMSNorm on queries and keys
Feed-forward	MoE and dense FFN layers alternate; the first FFN is dense; SwiGLU
MoE	LatentMoE: a shared down-projection to 3,072 before dispatch, 8 of 512 experts with softmax gating, expert FFN 10,240
Embeddings	`o200k_base` tokenizer (200,019 tokens), input and output embeddings tied
Global batch	134M tokens in every phase

The router sees the full hidden state, while the experts work on the compressed latent, so the all-to-all moves 3,072-wide vectors instead of 6,656-wide ones. The report says the query head count is rounded up to a multiple of 16 for tensor-parallel serving, and that the global layers use no position encoding because it is cheaper than RoPE at about the same quality.

2. Numerics#

Pre-training and mid-training precision#

The default type for weights and activations is BF16. The GEMMs split by pass:

GEMM	Precision
Forward	FP8 E4M3
Backward, data gradient	FP8 E5M2
Backward, weight gradient	BF16, with FP32 gradient accumulation

All FP8 operations use delayed scaling with a 1,024-step history of the absolute maximum.

FP32 is kept in the numerically sensitive places:

Activations: every pre-softmax activation (attention scores, MoE router logits, final output logits), the MoE combine, and the entire residual stream from the input embedding to the output.
Parameters and optimizer: the embedding, RMSNorm and router weights, the AdamW main parameters and momentum buffers, and all AdamW computation.
GEMMs and communication: the MoE routing GEMM, the final vocabulary GEMM, the data-parallel all-reduce, and the micro-batch gradient accumulation buffers.

Because each layer computes at lower precision than the FP32 residual stream, gradients have to be downcast in the backward pass. Stochastic rounding is applied whenever a gradient moves from higher to lower precision. Casts are fused into adjacent operations, such as RMSNorm, where possible.

Bitwise determinism#

For a fixed hardware topology, model configuration and software version, two pre-training runs produce bitwise identical models. The report accepts some loss of model FLOP utilization (MFU) for this.

GPU kernels. Scheduler-dependent atomics are replaced with accumulations in a fixed order. RMSNorm backpropagation uses a two-stage tiled reduction: one kernel writes per-tile partial sums, and a second reduces them in a fixed order. MoE top- $k$ selection uses a stable sort, so ties break the same way every time.
Collectives. NVLink SHARP is disabled so that all-reduce and reduce-scatter are deterministic. NCCL topologies are kept the same from run to run, so reductions inside a rack always accumulate in the same order.
Data and checkpoints. The micro-batch order is fixed for a given data-parallel layout, including across restarts. Checkpoints store the FP8 scaling history, dataloader progress and random-number generator state along with the weights and optimizer.

Determinism doubles as a regression test. YOLO’s tests run “golden” configurations and compare loss, gradient norms and weight checksums exactly. A change that alters floating-point behavior, such as a new kernel or driver, is allowed, but it updates the golden values, so the history records every expected numerical change.

3. Training: the YOLO framework#

All training phases use YOLO (“You Only Launch Once”), an in-house distributed training framework built on PyTorch.

Custom kernels#

YOLO’s kernels are written in Triton, CUDA, CuTe DSL and CUTLASS.

FP8 and MoE GEMMs. There are FP8 GEMMs for the forward and data-gradient passes with delayed scaling, and grouped GEMMs for MoE. Each GPU usually hosts 8 experts: 512 experts across expert parallelism (EP) of 64.
Fused quantization. One kernel quantizes both a tensor and its transpose, since the backward pass needs both, and writes the scale factors already swizzled.
Cluster Launch Control. The non-grouped quantization kernels use Cluster Launch Control (CLC), a Blackwell hardware feature, for dynamic load balancing.
Attention. FlashAttention-2’s deterministic mode was slow on GB200 and FlashAttention-4 (FA4) wasn’t available yet, so the team first wrote a custom block-sparse attention backend (a 1.06× step-time gain on v2). Switching to FA4’s deterministic kernels later gave 1.14× over that backend.
Expert encode. A PyTorch expert-encode kernel reached only 10% of HBM bandwidth. A Triton replacement reaches about 80% (a 1.03× end-to-end gain).
All-to-all. Custom CuTe DSL symmetric-memory kernels do device-initiated, variable-size all-to-all over NVLink.

Sharding and parallelism#

YOLO annotates each tensor with how it is split or replicated over an $n$ -dimensional mesh. The annotations are descriptive only. They never insert communication on their own, which avoids accidental synchronization points.

Per-tensor sharding. Each parameter’s sharding is chosen independently. Tensor and sequence parallelism (column- and row-parallel GEMMs) apply to the embedding, loss and attention weights, but not to the MLP or MoE weights.
Parallel folding. MoE and non-MoE layers can use different degrees of data parallelism.
Expert parallelism. EP is 64 with tensor parallelism 1, so all expert all-to-all traffic stays inside one rack’s NVL64 domain. Data parallelism runs across racks. DeepEP needs the all-to-all hidden dimension to be divisible by 512, so the LatentMoE latent dimension was set to 3,072 (6 × 512).
Data parallelism (ZeRO). YOLO has its own ZeRO stages 1 to 3. Parameters are grouped by sharding, dtype and ZeRO unit into contiguous buffers, which keeps the number of collectives low. The buffers are split into multi-dimensional chunks that stay compatible with checkpoint slicing. Parameters are always stored sharded, so the AdamW state is sharded at every stage; the stages differ only in when gathered buffers are freed. ZeRO-1 keeps gathered parameters until the last micro-batch, and ZeRO-3 frees them right after use. Pre-training and mid-training 1 use ZeRO-2; mid-training 2 (256K context) uses ZeRO-3 (FSDP).
ZeRO-3 to ZeRO-2 on v5. For the 1T-parameter v5 model, early runs used ZeRO-3, but the parameter all-gathers made the backward pass communication-bound. Activation offloading freed enough memory to go back to ZeRO-2, which restored the overlap between communication and compute.
Context parallelism. Long-context mid-training uses Ulysses-style context parallelism. An all-to-all moves the split from the sequence dimension to the attention-head dimension.
Pipeline parallelism is implemented too, though the report doesn’t say it was used for MAI-Base-1.
From capacity to dropless. Expert parallelism started GShard-style, with static shapes and a capacity factor of 2 (v2). From v3 on it is dropless, with variable-size all-to-all. The MoE layer also supports a global-batch load-balancing loss (the one used, aggregated across data-parallel workers and micro-batches), a loss-free balancer, and router replay for RL.

Memory and dropless MoE#

Checkpointing and offloading. YOLO combines whole-layer activation checkpointing with finer recomputation inside a layer. It also offloads activations asynchronously to pinned host memory on dedicated CUDA streams.
MoE pipelining. Local experts are split into groups, and the dispatch → compute → collect phases are pipelined across groups. All but the first dispatch and the last collect overlap with expert compute. The remaining exposed communication overlaps with the global load-balance computation and, optionally, the shared experts.
Static-memory dropless mode. Dropless routing never drops tokens, but an imbalanced batch can swing memory use and run out of memory. In this mode, each group runs several dispatch → compute → collect rounds, each capped at a fixed token count. In the backward pass, recomputation is done per expert and per round, so the imbalanced activations are never stored in HBM.

4. RL infrastructure: Rocket and SGLang#

RL runs on Rocket, an in-house asynchronous RL framework. It uses YOLO as the learner and SGLang with the SGLang router for inference. The team built it because no open-source RL framework they looked at handled asynchronous RL across several thousand GPUs.

Components#

Rocket has one controller, a pool of problem workers, a pool of rollout workers, and the SGLang router and servers. The controller and each worker are single Python processes running as Ray actors.

Controller. Loads tasks, sends them to problem workers, filters the graded rollouts that come back, and batches them for the learner. Large runs use the off-policy controller; the on-policy one is for small experiments and debugging.
Problem worker. Handles one task. It first sends 16 rollout requests (early exit) and drops the task unless the pass rate is between 0.05 and 0.8. It then sends 128 more and keeps the task only if the full pass rate is between 0.1 and 0.8. Last, it computes normalized advantages, applying length penalties or reward normalization where the task calls for them. It retries rollouts whose worker failed.
Rollout worker. Generates one rollout: it sends the prompt to an inference server, runs any tool calls in the reply, and loops until the model gives a final answer or hits a step or time limit. Graders that judge one rollout alone run here; graders that compare rollouts, such as pairwise LLM judges, run on the problem worker.
Length curriculum. Rollouts start capped at 8K tokens, and the cap doubles until it reaches 128K. Short caps make the early, low-accuracy phase much cheaper to serve.

GPU allocation#

Inference takes most of the GPUs. The report puts the inference-to-learner ratio as high as 5:1. In the largest job, 4,096 of 4,864 GB300s served SGLang and 768 ran the learner, which works out to 5.3:1.

Inference by workload#

Single-turn (for example STEM). Prompts are short, but generations with thinking reach 128K tokens, so KV cache memory is the bottleneck.
- Expert parallelism and data parallelism for attention shrink the KV cache and weight footprint per GPU.
- Data parallelism on the MLP layers, DeepEP, and EPLB (DeepSeek’s expert-parallel load balancer) cut communication.
- Prefix caching is turned off, so that tokens in the sliding-window attention layers can be fully evicted during long generations. The model pairs 5 sliding-window layers (window 512) with each global layer.
Multi-turn (for example agentic SWE and tool use). Prompts grow across turns while each generation is short, so the load is mostly prefill. Here prefix caching does the work: the hit rate reaches 97 to 98% on production RL runs. It helps that every turn’s context is a strict prefix of the next.

Inference stability#

With thousands of inference GPUs, replicas crash, hang or slow down regularly, and the job has to survive it. The report describes three layers:

Replica. Each SGLang server runs a watchdog that probes its own generation endpoint and watches scheduler memory. An unhealthy replica restarts itself so the orchestrator can replace it.
Router. The SGLang router acts as a circuit breaker. When backends go unhealthy, it stops admitting requests and runs a multi-stage probe before reopening. It also periodically rediscovers replicas that silently fell out of its active set. Per-replica flow control keeps one slow replica from building a backlog.
Job. A liveness monitor counts live actors of each kind (inference replicas, routers, rollout workers, learner ranks) and restarts the job if any kind falls below its threshold. A separate step-progress watchdog catches the case where everything is alive but training has stopped advancing.

A request that fails mid-flight is retried on another replica, so a replica failure costs throughput but no rollouts.

Sandboxes and environments#

Agentic rollouts run tools in the Sandbox Execution Environment (SEE), which creates a fresh container for each task and destroys it afterwards. Containers have no network by default, which keeps episodes deterministic and free of rate limits and transient failures. Tasks that need the network, for example to install packages, go through caching proxies and domain allowlists.

The SWE environments are built by a separate pipeline on a Ray cluster:

Two pools. A main pool dispatches work, handles Lance I/O and tracks pipeline state. A builder pool runs the container builds. Builds fail often (full disks, out of memory, hangs), so isolating them keeps those failures away from orchestration. Builder pods also run untrusted Dockerfiles from open-source repositories, so they hold no credentials for the pipeline’s storage or internal services.
Builder pods. Each pod has two containers on shared NVMe: a Ray worker with rootless podman, which loads images and runs grading containers, and a rootless BuildKit sidecar. BuildKit writes OCI archives to the NVMe volume and podman loads them, so images never cross the network between build and grading.
Layer caching. Problems from the same repository go to the same persistent actor, so shared layers (system packages, pip installs, dependency builds) are built once.
Throughput. The bottleneck is LLM tokens, not CPU or disk: an LLM agent writes the Dockerfiles. With 10,000 CPU cores and 12 million LLM tokens per minute (83% prompt cache hit rate), the pipeline produces about 20 environments per minute that pass grading. The appendix also describes the cluster as about 30,000 CPU cores; it doesn’t reconcile the two figures.

Self-distillation as a numerics reset#

Early in the RL work, before the fixes in the next subsection, small inference–learner differences sometimes built up until a climb diverged. Resuming from an earlier checkpoint often didn’t help, because the instability was already in the weights many steps before the collapse. Instead, the team fine-tuned a mid-trained checkpoint on rollouts from the failed climb (packed 128K sequences, global batch 2,048) and restarted RL from that. The same step carries progress over to a new base checkpoint. The report says collapses needing this became rarer as the infrastructure improved.

Closing the learner–inference numerics gap#

YOLO and SGLang use different kernels, scheduling and parallelism. Small per-token log-probability differences compound over long rollouts and can destabilize the importance-sampling correction in off-policy RL. Three mechanisms keep the gap small:

BF16 on both sides. The learner and SGLang both run BF16 during RL. The report says this gave smaller gaps and more stable training than the lower-precision alternatives it tried.
MoE routing replay. The expert choices made during inference are recorded and replayed in the learner’s pass.
Top- $p$ mask replay. Rollouts are sampled with $p = 0.97$ . The truncation mask is stored and replayed in the learner’s forward and backward pass, setting the logits of excluded tokens to $-\infty$ before the softmax. Without it, backpropagating through tokens outside the nucleus caused divergence within a few steps. A larger $p$ explores more but costs more to transfer the masks.

Asynchronous weight transfer#

Cadence. The learner pushes fresh weights to the inference fleet every 5 gradient steps. Rollouts generated by a policy more than 8 updates (40 gradient steps) old are discarded.
Transfer plan. The learner and inference sides shard tensors differently (FSDP, pipeline stages, data-parallel attention, layouts), and parameters may differ in precision or quantization. At job start, Rocket computes a transfer plan that intersects the two layouts and records one entry per overlapping sub-shard: source rank, destination rank, byte range, and any dtype cast or layout permutation. Each transform runs on whichever side it is cheaper.
Slices only. Each learner rank sends only the slices its paired inference rank needs, without assembling full tensors. Packing, transfer and unpacking are pipelined.
Fleet-size independence. The plan is written for one learner and one inference server. At runtime each source–destination pair expands to every live replica of the destination, so the plan survives replicas joining and leaving.
Parallel groups. Data parallelism replicates parameters, so one data-parallel group already holds every weight, and the transfer uses only a subset of groups. For example, a 36-server fleet split four ways runs four independent 9-server transfers.

5. Performance#

MFU is measured against the GB200’s dense FP16/BF16 tensor-core peak, $2.5 \times 10^{15}$ FLOPS per GPU. The count excludes memory-bound operations and activation recomputation. Since some GEMMs run in FP8 but are scored against the BF16 peak, these MFU numbers aren’t directly comparable to an all-BF16 run.

Version	Change	MFU before → after optimization
v2 (4K GB200)	first GB200 run, EP 64 inside NVL64	18% → 22%
v3	dropless MoE	about 22%
v4 (8K GB200)	512 experts, top-8, LatentMoE	16% → 20%
v5 (MAI-Base-1)	35B active, 1T total	18% → 20%

More than 20 kernel and infrastructure optimizations kept each generation above 20% MFU on up to 8,192 GPUs. The report names the main ones:

v2 (18% → 22%). Profiling found three bottlenecks: data-parallel communication, deterministic attention and expert routing. GPUDirect RDMA improved communication overlap (about 1.1×). The custom block-sparse attention gave 1.06×. ZeRO-2 cut gradient memory with no more communication than ZeRO-1, since there is a single micro-batch, which let some activation recomputation be turned off (1.03×). The Triton expert-encode kernel gave another 1.03×.
v3 (about 22%). Dropless routing removed capacity padding, which cut both all-to-all volume and expert GEMM work. In exchange, token counts and tensor shapes have to be exchanged at runtime. Overlapping that exchange with other work and moving synchronization points onto dedicated streams kept MFU level.
v4 (16% → 20%). Experts went from 192 to 512, routing from top-4 to top-8, and the job from 4,096 to 8,192 GPUs. LatentMoE shrank the expert GEMMs, which made CPU launch overhead and kernel efficiency matter more. FA4’s deterministic kernels (1.14× over the block-sparse backend) and work on CPU overhead and batching recovered the loss.
v5 (18% → 20%). The jump to 35B active and 1T total parameters forced ZeRO-3, which made the backward pass communication-bound, until activation offloading allowed the return to ZeRO-2 (section 3).

Version	Active / total	Layers	Hidden	Experts (top-k)	Capacity factor
v2	23B / 600B	54	4,352	192 (4)	2
v3	23B / 600B	54	4,352	192 (4)	dropless
v4	23B / 611B	66	5,632	512 (8), LatentMoE	dropless
v5 (MAI-Base-1)	35B / 1T	78	6,656	512 (8), LatentMoE	dropless

6. Reliability and fault tolerance#

Goodput. The report defines goodput as ideal training time divided by actual wall-clock time. The 8K-GPU MAI-Base-1 pre-training run reached 90.0%, with 51 hours of total overhead. The named categories are MFU drops (18 hours, 35%), non-stepping time (14 hours, 27%), and recomputation after checkpoint fallback (6.5 hours; the report says 15%, though 6.5 of 51 is 13%). These sum to 38.5 hours, so about 12.5 hours fall in categories the report doesn’t break out. By my arithmetic, 51 hours at 10% overhead implies a run of about 510 hours, or 21 days.
Checkpointing. YOLO rewrote the metadata and serialization paths of PyTorch Distributed Checkpointing (DCP) and caches precomputed save plans. Together these cut CPU memory overhead and save time by more than 10×. Tensors are copied from device to host inside the training processes, and a separate process writes them to the datacenter’s local Azure Blob Storage while training continues. Cross-rank commit uses a custom coordination layer on a shared key-value store in place of the CPU collective backend. A new checkpoint starts as soon as the previous one finishes, with at most one in flight. Local Blob Storage is fast enough that there are no node-local caches or cross-rank state replicas. On load, replicated state is read from storage once and broadcast over NCCL, which avoids overloading a single storage server.
Recovery. Jobs restart in place using Ray actors with hot standbys, so pods aren’t recreated. On GB200, the runtime clears OS caches during actor restarts. Without that, GPUs holding memory from the previous run caused out-of-memory crash loops. After a restart, determinism is checked by comparing the recomputed steps’ loss with the logged history.
Certification. A node must pass three stages before it enters the Kubernetes/Kueue pool. First, isolated stress tests of GPUs, CPUs, HCAs, NVLink and memory, looking for high ECC rates, throttling, clock violations and link flaps. Second, rack-level NCCL stress tests, which found several GB200 and GB300 racks with fewer than 16 healthy nodes; those racks were kept out of large runs. Third, cross-rack InfiniBand tests across racks, rails and leaf groups to validate the spine layer, run on demand. Repaired nodes go through certification again.
Node lifecycle. A new or repaired node enters Init, then Certifying. A pass makes it Available; a failure makes it Impaired. Runtime signals (node-problem-detector conditions, XID errors, ECC thresholds, NVLink degradation, InfiniBand flaps, storage faults) can also move an Available node to Impaired. Transient problems get automatic remediation, such as a reboot or soft drain. Hardware problems go to vendor or datacenter maintenance, and the node returns to Init afterwards. Nodes that can’t be fixed are decommissioned.

7. Scheduling and observability#

The scheduling stack is layered:

Kubernetes holds cluster state, usually one Kubernetes cluster per datacenter building. Large jobs can span clusters, and pods use host networking to avoid overlay overhead.
Kueue handles priority, quota, admission, preemption and topology-aware placement, trying to land large jobs in compact regions of the high-bandwidth domains.
MAI’s own controllers keep reservations, rack topology labels (set at node start from the physical host) and readiness gates up to date.
Ray runs the job. MAI’s drivers turn the admitted topology into actor placement, communication groups and NCCL clique configuration.
A fleet-wide control plane gives one view across clusters and scheduler backends.

Rack fragmentation is the recurring problem: small jobs from many queues can fill racks so that a later large job can’t find contiguous capacity. The fix is soft rack reservations. Each queue has preferred racks, other queues can borrow them while idle, and Kueue’s reclaimWithinCohort lets the owner take them back.

For observability, GPU health is tracked through XIDs, ECC, thermals, power, clock throttling, NVLink state and bit-error rate, chip-to-chip links, InfiniBand device state, NVMe health, PCIe errors and driver state. These become Kubernetes node conditions, and custom triage and drain controllers act on them, so the scheduler sees health directly. Telemetry goes to Datadog (near-real-time metrics and logs), Azure Managed Prometheus (time series) and Azure Data Explorer (long retention). Alerts are layered the same way across Prometheus rules, Datadog monitors, KQL checks and Azure Monitor.