Ai2 Olmo-core 3: The Secret Behind 2.7x MoE Training Throughput

A System-Level Rewrite Targeting Sparse Model Training Bottlenecks

The Allen Institute for AI (Ai2) has released Olmo-core 3, its open-source training framework for the Olmo family of language models. The release is available under the Apache 2.0 license and comes with a full technical report. The core change points in one direction: systematically dismantling the overhead in Mixture-of-Experts (MoE) model training that inflates with total parameter count without delivering a proportional amount of useful computation.

According to the technical report, the runnable configurations documented in Olmo-core 3 range from a 12.9B parameter model on 16 GPUs all the way to a 1.2T parameter model on 512 GPUs, complete with official pretraining scripts and checkpoint conversion tools for exporting to Hugging Face format.

The Capacity Tax: A Structural Burden of Sparse Models

The technical report gives this class of overhead a clear name: the capacity tax. It refers to costs incurred by total stored parameter count — weight storage, gradient synchronization, optimizer state updates — even when the number of parameters actually activated per token remains constant. These costs still scale linearly with total parameter count.

Dense models don't run into this problem, because more parameters means more effective computation — cost and benefit grow in lockstep. Sparse models break that coupling. Early Fully Sharded Data Parallel (FSDP) implementations had to gather and reshard weights on every micro-batch, with the result that increasing the number of experts at a fixed activation size necessarily degraded throughput.

Olmo-core 3's answer is to switch to Distributed Data Parallel (DDP): each rank holds its own model partition throughout the entire accumulation window, gradients are synchronized only once per optimizer step, and a distributed optimizer shards the FP32 master weights and optimizer states. On top of that, expert parallelism spreads routed experts across ranks, writing only the selected token rows directly into the target rank's buffer via NVSHMEM. Layers that don't fit are covered by pipeline parallelism.

Mixture-of-Experts (MoE) models are built around splitting the feed-forward network into multiple parallel "expert" sub-networks, with each token activating only a few of them (say, 4 out of 128) via a routing function. This allows the total parameter count to balloon enormously while the actual computation per token (i.e., activated parameters) stays constant. The problem is that weight storage, gradient synchronization, and optimizer state are all computed based on total parameter count rather than activated parameters — so the larger the total, the higher this computation-independent overhead becomes. This is precisely where the term "capacity tax" comes from. Sparse models are therefore fundamentally different from dense models: dense models scale parameters and effective computation together, while sparse models decouple the two, shifting the burden onto the system implementation layer.

The key difference between Fully Sharded Data Parallel (FSDP) and Distributed Data Parallel (DDP) lies in when weights are synchronized. FSDP shards model parameters across GPUs but must gather the weights back before every forward and backward pass, then reshard them afterward. For MoE, the more micro-batches there are, the more frequently this gather/reshard round-trip occurs, making it a throughput bottleneck. DDP, by contrast, lets each rank hold its own model partition across the entire gradient accumulation window, synchronizing gradients only once per optimizer step, with a distributed optimizer sharding the FP32 master weights and optimizer states — thereby amortizing synchronization overhead.

Measured Data on Expert Pool Scaling

The core question this rewrite sets out to answer is: how large can the expert pool grow without eating into throughput?

Ai2 provides two key measurements. In the first, the expert pool scales from 8 to 128 experts, with each token selecting only 4 experts, keeping activated parameters constant at roughly 3.2B. Total parameter capacity grows from 4.6B to 47B — a tenfold increase — while training throughput drops by less than 5%.

The second comparison runs on 8 NVIDIA B300 GPUs: a 47B-parameter MoE model reaches approximately 52,000 tokens per second per GPU on the new framework, versus only about 19,400 on the old implementation — a throughput improvement of roughly 2.7x. Ai2 labels this run a "preliminary result."

Routing Is Not Just a Model Problem — It's a Systems Problem

The report treats routing as an engineering interface rather than a modeling detail. One section digs into what the load-balancing loss function is actually optimizing, and names a failure mode it calls "Token Gerrymandering": the loss value improves while imbalance across experts actually worsens. Any team that has tuned routing uniformity has run into similar results.

Another finding is worth noting: the report explicitly documents optimizations the team tried and then abandoned. For example, overlapping communication with computation — which intuitively should hide communication overhead — actually slowed the overlapped kernels by more than the overhead it hid, because the two were competing for the same GPU resources.

Two engineering details are directly valuable for teams reading the code. Routing metadata is kept on the GPU, eliminating the host round-trips the old path made for all-to-all split sizes and expert matrix-multiplication group sizes. Device-side scheduled grouped GEMM handles the uneven per-expert shapes that routing produces, and the report also measures how this compares against padding all expert batches to a uniform capacity.

NVSHMEM is NVIDIA's Partitioned Global Address Space (PGAS) programming model library, which allows GPU kernels to read and write remote GPU memory directly without first going through host-side MPI communication. This is critical for MoE expert parallelism: routing determines which tokens need to go to the rank hosting a given expert, so Olmo-core 3 writes the selected token rows directly into the target rank's buffer via NVSHMEM, avoiding traditional host round-trips and explicit all-to-all calls. Similarly, grouped GEMM handles the problem of uneven token counts received by each expert after MoE routing — it packs small matrix multiplications of different shapes into a single kernel for batched execution, and device-side scheduling further eliminates the extra computation and memory footprint of padding for shape alignment.

Boundary Conditions of the Throughput Numbers

The report is clear about the conditions under which these numbers should be interpreted, and they warrant caution when cited. The measured configurations differ in batch size, parallelism strategy, precision, and recomputation strategy, so the authors describe them as a set of "operating points" rather than a controlled scaling curve. Each headline number is a system reference value measured under random routing conditions.

Precision choices directly affect economics. MXFP8 — an eight-bit format with shared block scaling — reduces GEMM time, saved activations, and dispatch payloads, while retaining an authoritative copy of the FP32 master weights. Under a 1.2T-parameter configuration, combined with MXFP8 and layer-wise recomputation, it reached 858 useful-model TFLOP/s per GPU. An optional DeepEP v2 backend touched 2.38T parameters in an experimental capacity test.

MXFP8 is an eight-bit floating-point micro-scaling format defined by the Open Compute Project. Its distinguishing feature is that every 32 elements share a single scaling factor (block scaling), rather than a single scale for the entire tensor. Compared with conventional per-tensor FP8, MXFP8 maintains numerical stability while providing finer-grained compression for GEMM computation, activation savings, and expert dispatch payloads. Olmo-core 3 uses it to accelerate matrix multiplication and reduce memory footprint, while keeping an authoritative FP32 master weight copy for optimizer updates so that precision doesn't degrade from quantization. Layer-wise recomputation, meanwhile, discards some intermediate activations during backpropagation and recomputes them when needed — trading memory for computation, a common technique for fitting large models into limited GPU memory.

Where the Costs Are

The throughput gains do not come for free. Activation recomputation reduces activation memory by nearly two-thirds and peak memory by about a quarter, at the cost of roughly one-fifth of throughput. This is the trade-off large configurations make to fit into device memory.

Topology-agnostic checkpointing preserves portability across different parallel layouts — training jobs can resume on a different number of nodes without rewriting the checkpoint format. For teams maintaining their own GPU clusters, this reduces operational friction rather than computational cost.

The Olmo-core 3 release itself is a signal: competition in MoE training is shifting from "model architecture design" to "system implementation efficiency." When the expert pool can grow tenfold with almost no change in throughput, the path to scaling sparse models is no longer blocked by infrastructure.

Topology-agnostic checkpointing means the checkpoint format is not tied to the parallel layout used during training — for example, a model originally trained on 512 GPUs with a specific tensor-parallel and pipeline-parallel split can be resumed directly under a different node count and parallelism strategy. Traditional checkpoints often record how parameters are distributed across specific ranks, so changing the layout requires rewriting the format or doing offline conversion. Topology-agnostic schemes instead store parameters by logical structure and dynamically map them to the new layout on resume. This lets teams quickly restart training jobs after cluster failures, scaling changes, or hardware swaps — reducing operational friction rather than the computational cost itself.

One-line Summary

Ai2's open-source Olmo-core 3 eliminates the capacity tax in MoE training through a system-level rewrite, enabling a tenfold expert pool expansion with less than 5% throughput loss.

Paragraph Summary

The Allen Institute for AI has released Olmo-core 3 under the Apache 2.0 license, an open-source framework that systematically rewrites the bottlenecks of Mixture-of-Experts (MoE) model training. Its core goal is to eliminate the "capacity tax" — overhead that inflates with total parameter count without delivering a proportional amount of effective computation. The approach replaces Fully Sharded Data Parallel (FSDP) with Distributed Data Parallel (DDP), combined with expert parallelism, direct NVSHMEM writes, a distributed optimizer, and pipeline parallelism. Measurements show that scaling the expert pool from 8 to 128 and total parameters tenfold yields less than 5% throughput loss; a 47B MoE reaches roughly 52,000 tokens per second per GPU on 8 B300s, about 2.7x the old implementation. The report also discloses a routing load-balancing failure mode, an abandoned communication-computation overlap approach, and the boundary conditions of its key numbers.