Ai2 releases Olmo-core 3 under Apache 2.0 with 2.7x MoE training throughput gains

What Ai2 shipped

On 1 October 2026 the Allen Institute for AI (Ai2) published version 3 of Olmo-core, the framework it uses to train its Olmo language models, as open source under the Apache 2.0 licence together with a technical report. According to a dev.to write-up of the release, the package installs from PyPI as ai2-olmo-core and ships with official pretraining recipes for the Olmo 3 7B and 32B models, plus checkpoint conversion tools that export to Hugging Face format.

The capacity problem it targets

The rewrite centres on the economics of sparse models. A mixture-of-experts (MoE) network stores many expert sub-networks but routes each token through only a few of them, so the compute per token stays small while the stored parameter count balloons. Ai2's report, as described on dev.to, separates training costs into two families: a "capacity tax" that scales with stored parameters, covering weight storage, gradient synchronisation and optimizer updates, and a "runtime tax" tied to routing and execution. Dense models absorb the first cost naturally because more parameters also mean more useful work. Sparse models break that link.

Under the previous implementation, weights were gathered and resharded around every microbatch, so adding experts at a fixed active size dragged throughput down even though the useful work per token barely changed. That mismatch motivated the rebuild.

How the new stack works

Olmo-core 3 swaps fully sharded data parallelism for distributed data parallelism. Each rank keeps its partition of the model resident for the whole accumulation window and synchronises gradients once per optimizer step, while a distributed optimizer shards the FP32 main weights and optimizer state. Expert parallelism then spreads the routed experts across ranks and sends only the selected token rows to whichever rank holds the matching expert, using NVSHMEM to write each row directly into its destination buffer. Pipeline parallelism covers layers that will not fit on a single rank.

Two implementation details stand out for anyone reading the code. Routing metadata now stays on the GPU, removing a round trip to host memory that the older path required for all-to-all split sizes and expert group sizes. Device-scheduled grouped GEMMs execute the uneven per-expert shapes that routing produces; the report measures where this beats padding every expert batch to a common capacity, and where it does not.

The numbers Ai2 reports

In one benchmark, Ai2 grew the expert pool from 8 to 128 experts while selecting four per token, holding active parameters near 3.2 billion. Total capacity rose from 4.6 billion to 47 billion parameters, and throughput fell by less than 5%. In a second run that Ai2 describes as preliminary, a 47-billion-parameter MoE model processed roughly 52,000 tokens per second per GPU on eight NVIDIA B300 GPUs with the new stack, against about 19,400 with the implementation it replaces, roughly 2.7 times the throughput.

The report's measured operating points span a 12.9-billion-parameter configuration on 16 GPUs up to a 1.2-trillion-parameter configuration on 512 GPUs, where it reached 858 useful-model TFLOP/s/GPU using MXFP8 precision and per-layer recompute. An experimental capacity test on the optional DeepEP v2 backend reached 2.38 trillion parameters. Ai2 cautions that the configurations differ in batch size, parallelism, precision and recompute policy, so the figures are operating points rather than a controlled scaling curve, and each headline rate is a systems reference measured under random routing.

There are explicit trade-offs. Activation recompute removes nearly two-thirds of activation memory and about a quarter of peak memory, at a cost of roughly a fifth of throughput. Topology-agnostic checkpoints let a run resume on a different node count without rewriting its checkpoint.

What the team rejected

The report also documents the approaches the team tried and did not adopt, arguably its most useful section for other practitioners. Overlapping communication with computation slowed the overlapped kernels by more than the overlap hid, because both activities competed for the same GPU resources. On the modelling side, the authors name a failure pattern they call "Token Gerrymandering", in which the load-balancing loss improves while imbalance between experts grows.

Why it matters

Production-scale training code is rarely published at this level of detail, and Olmo-core 3 gives practitioners a complete, reproducible large-scale training stack under a permissive licence, from pretraining scripts to distributed optimizer internals. The headline lesson is structural: as expert counts grow, parameter bookkeeping rather than arithmetic dominates the cost, and keeping experts resident while moving only token rows claws much of that back. Just as valuable, the documented negative results spare other teams from repeating experiments that sound promising in theory.