Ai2's Olmo-core 3 Keeps Every Expert on the GPU and Throughput Nearly Flat

TL;DR

• Olmo-core 3 is the open training framework the Allen Institute for AI builds its Olmo language models with, released as source code under the Apache 2.0 licence with a technical report beside it.
• The release appeared on 1 October 2026 and documents measured operating points from a 12.9-billion-parameter model on 16 GPUs up to a 1.2-trillion-parameter model on 512 GPUs.
• The central change swaps the earlier fully sharded data parallelism for distributed data parallelism, so each rank keeps its model partition resident and selected token rows travel to the expert that holds them, which removes the cost the report calls the capacity tax that follows stored parameters.
• A 47-billion-parameter mixture-of-experts model reached about 52,000 tokens per second per GPU on eight NVIDIA B300 GPUs against about 19,400 with the earlier implementation, roughly 2.7 times the throughput.
• Growing the expert pool from 8 to 128 experts, with active parameters held near 3.2 billion, cost under 5% of throughput, and the report labels every headline rate a systems reference measured under random routing.
• The report also names the optimisations the team tried and did not adopt, including overlapping communication with computation, which slowed the kernels it was meant to hide.

What the framework changes

Olmo-core is Ai2's training framework, and version 3 rebuilds the part that trains mixture-of-experts models. The upgrade ships as source code under the Apache 2.0 licence, installable from PyPI as ai2-olmo-core, with official pretraining scripts for the Olmo 3 7B and 32B models and checkpoint conversion tools that export to Hugging Face format.

The motivation sits in the arithmetic of sparse models. A mixture-of-experts model stores many expert networks and sends each token through only a few of them, so the compute per token stays small while the stored model grows. Weight storage, gradient synchronisation and optimizer updates keep tracking the total parameter count, and the earlier implementation gathered and resharded weights around every microbatch. Adding experts at a fixed active size therefore reduced throughput even though the useful work per token barely moved.

The capacity tax behind the rewrite

The technical report gives that mismatch a name. Its authors separate two cost families: the capacity tax, which follows stored parameters, and the runtime tax of sparse routing and execution, which follows the ways a distributed system spends time. Dense models carry the first without noticing, because more parameters also mean more useful computation. Sparse models break that link.

The new design starts from distributed data parallelism. Every rank holds its partition of the model for the whole accumulation window and synchronises gradients once per optimizer step, with a distributed optimizer sharding the FP32 main weights and optimizer state. Expert parallelism then shards the routed experts across ranks and dispatches only the selected token rows to the rank that stores their expert, using NVSHMEM to write each row straight into its destination buffer. Pipeline parallelism covers the layers that will not fit on one rank.

Two smaller decisions matter for a team reading the code. Routing metadata stays on the GPU, which removes the round trip to the host that the older path made for both the all-to-all split sizes and the expert matrix-multiplication group sizes. Device-scheduled grouped GEMMs then execute the uneven per-expert shapes routing produces, and the report measures where that beats padding every expert batch to a common capacity and where it does not.

An expert pool that grows without eating throughput

The headline measurement answers the question the rewrite was built for. In one benchmark Ai2 grew the expert pool from 8 to 128 experts while selecting only four experts per token, holding active parameters per token near 3.2 billion. Total parameter capacity went from 4.6 billion to 47 billion, and training throughput fell by less than 5%.

A second benchmark compared the new stack with the implementation it replaces on eight NVIDIA B300 GPUs. A 47-billion-parameter mixture of experts processed 52,000 tokens per second per GPU on the new stack against 19,400 on the old one, which is about 2.7 times the throughput. Ai2 describes that run as preliminary.

Routing as a systems problem

The report treats routing as an engineering surface rather than a modelling detail. One section works through what the load-balancing loss actually optimises, and it names a failure pattern the authors call Token Gerrymandering, where the loss improves while the imbalance between experts grows. Any team that has tuned a router for evenness has met a version of that outcome.

Another finding cuts against a common intuition. Overlapping communication with computation slowed the overlapped kernels by more than the overlap hid, because the two activities competed for the same GPU resources. The report documents the approaches the team did not adopt alongside the ones it did, which is the part of a systems paper that saves other people the same experiments.

Reading the numbers in the report

The report states its own terms clearly, and they matter for anyone quoting the figures. Its measured operating points run from a 12.9-billion-parameter configuration on 16 GPUs to a 1.2-trillion-parameter configuration on 512 GPUs, reaching 858 useful-model TFLOP/s/GPU at 1.2 trillion with MXFP8 and per-layer recompute. An experimental capacity test on the optional DeepEP v2 backend reaches 2.38 trillion parameters.

The configurations differ in batch size, parallelism, precision and recompute policy, so the authors describe the rates as a set of operating points rather than a controlled scaling curve, and each headline number is a systems reference measured under random routing. The table states the batch, topology, precision and recompute policy needed to interpret each figure. Precision details change the economics as well: MXFP8, an eight-bit format with shared block scales, cuts GEMM time, saved activations and dispatch payloads while one FP32 main weight stays authoritative.

Where the throughput gains stop

Two costs show up in the same document. Activation recompute removes nearly two thirds of activation memory and about a quarter of peak memory, and it takes about a fifth of throughput in return, which is the trade the larger recipes make so they fit in device memory. Topology-agnostic checkpoints keep a run portable between parallel layouts, so a training job can resume on a different node count without rewriting its checkpoint.

For a team running models on its own hardware, including the pattern behind a self-hosted AI agent platform, the practical read is a budget question. Training a large mixture of experts still needs a GPU cluster, and the framework removes part of the waste that arrives with scale. Our note on benchmarking LLM inference with a real load client covers the same discipline on the serving side, where a published number needs its load conditions beside it.

Sources

Source: Introducing Olmo-core 3, open and scalable training infrastructure for large MoEs — Ai2, published 2026-10-01 — https://huggingface.co/blog/allenai/olmocore3, retrieved 2026-10-10 (redesigned open mixture-of-experts training system designed to scale into the trillion-parameter range; expert pool grown from 8 to 128 experts with four experts selected per token and active parameters near 3.2B; total parameter capacity from 4.6B to 47B with throughput falling less than 5%; the same infrastructure benchmarked past one trillion total parameters; the earlier implementation's fully sharded data parallelism against the new distributed data parallelism that keeps experts resident on GPUs; 52,000 tokens per second per GPU against 19,400 on eight NVIDIA B300 GPUs, about 2.7x).

Source: Supercharging Olmo-core for Efficient and Scalable MoE Training — Ai2 technical report, October 2026 — https://allenai.org/papers/olmocore3, retrieved 2026-10-10 (DDP architecture with expert parallelism, pipeline parallelism, a distributed optimizer and activation recompute; NVSHMEM-based rowwise expert parallelism; GPU-resident routing metadata and device-scheduled grouped GEMMs; MXFP8; topology-agnostic checkpoints; measured configurations from 12.9 billion parameters on 16 GPUs to 1.2 trillion on 512 GPUs reaching 858 useful-model TFLOP/s/GPU, with an experimental capacity test reaching 2.38 trillion; activation recompute removing nearly two thirds of activation memory and about a quarter of peak memory at a cost of about a fifth of throughput; the capacity tax and the runtime tax; the failure pattern the authors call Token Gerrymandering; the finding that overlapping communication with computation can slow the overlapped kernels by more than it hides; the statement that every headline rate is a systems reference measured under random routing; the use of AI tools in developing the codebase and report).

Source: Ai2 Releases OlmoCore 3 — Tech Times, published 2026-10-02 — https://techtimes.com/articles/328470/20261002/ai2-releases-olmocore-3-open-moe-training-stack-delivers-27x-throughput-gain.htm, retrieved 2026-10-10 (Apache 2.0 licence, PyPI package ai2-olmo-core, official pretraining scripts for Olmo 3 7B and 32B, checkpoint conversion to Hugging Face format, GPU cluster access as the primary constraint).

Source: Ai2 — Olmo-core 3 launch post and technical report, published 1 October 2026, retrieved 2026-10-10. Figures from the launch post and the report's own measured operating points; images are Ai2's published charts and a capture of the launch post, plus two Netics editorial diagrams.