DEV Community

Cover image for Ai2's Olmo-core 3 Keeps Every Expert on the GPU and Throughput Nearly Flat
Taha Echakiri
Taha Echakiri

Posted on Originally published at blog.neticslabs.com

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.

Netics editorial definition card for the capacity tax, the share of training cost that follows total stored parameters even when active parameters per token stay fixed, with the DDP and expert-parallelism approach Olmo-core 3 uses in place of it.

Netics editorial definition card: the capacity tax described in Ai2's technical report, and the routing arrangement Olmo-core 3 uses in its place.

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.

Expert-pool scaling chart from Ai2's Olmo-core 3 launch post, showing the expert pool growing from 8 to 128 experts while training throughput falls by less than 5%.

Ai2's own expert-pool scaling chart from the Olmo-core 3 launch post (huggingface.co/blog/allenai/olmocore3, published 2026-10-01, retrieved 2026-10-10): the expert pool on one axis and training throughput on the other.

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.

Screenshot of Ai2's Olmo-core 3 launch post as published on Hugging Face, showing the title 'Introducing Olmo-core 3: Open, scalable training infrastructure for large MoEs', the October 1 2026 publication date and the technical report, code and demo links.

Capture of Ai2's official Olmo-core 3 launch post on Hugging Face (huggingface.co/blog/allenai/olmocore3, published 2026-10-01, captured 2026-10-10): the post that carries the charts and links the technical report.

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.

Training throughput comparison chart from Ai2's Olmo-core 3 launch post, showing the new stack at 52,000 tokens per second per GPU against 19,400 for the earlier implementation.

Ai2's own throughput comparison chart from the same launch post (retrieved 2026-10-10): the two implementations measured on the same eight B300 GPUs.

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.


Originally published on the Netics blog.

Top comments (0)