Ai2 has released Olmo-core 3, an open training framework for large mixture-of-experts models. The redesign keeps experts on GPUs and sends token data to them, reducing the repeated movement of model weights. Builders can use the code to train their own models and experiment with routing and how they divide work across a cluster.
Ai2 reports about 2.7 times the throughput of its earlier implementation in a preliminary test of a 47-billion-parameter model on eight NVIDIA B300 GPUs. The new stack processed 52,000 tokens per second per GPU, compared with 19,400 for the earlier version.
More experts with similar active computation
A mixture-of-experts model selects a subset of its experts for each token. That lets it add parameters without using all of them on every input. The full expert pool still needs memory, and GPUs must exchange data to route tokens to the right experts.
In Ai2’s expert-scaling benchmark, the pool grew from eight to 128 experts while each token still selected four. Total parameters rose from 4.6 billion to 47 billion. Active parameters stayed near 3.2 billion per token, and throughput fell by less than 5%.
The published chart specifies eight B300 GPUs, BF16 precision and random routing. Those conditions matter when comparing the result with a model that learns which experts to select.
What changed in the training code
The earlier implementation used fully sharded data parallelism, gathering and resharding weights for each small batch. Olmo-core 3 uses a system based on distributed data parallelism that keeps experts resident on GPUs.
Expert parallelism divides the expert pool across GPUs. Pipeline parallelism divides the model’s layers across GPU groups. A distributed optimizer divides the extra state needed to update parameters, reducing the amount each GPU stores.
The stack also places routed data directly into expert input buffers and keeps routing metadata on GPUs. Grouped matrix multiplication combines small expert computations into work the GPUs can execute more efficiently.
Ai2 tested lower-precision MXFP8 against BF16 on four B300 GPUs with work distributed uniformly across experts. Enabling MXFP8 where it helped most raised end-to-end throughput by about 21% and reduced peak active memory from 103 GiB to 95 GiB.
What the trillion-parameter tests establish
Ai2 benchmarked a 1.2-trillion-parameter configuration on 512 B300 GPUs, with 58.36 billion parameters active per token. Its highest observed throughput was 858 TFLOP/s per GPU. The tests used random routing to measure system performance rather than trained-model quality.
A separate experiment with DeepEP v2 reached 2.38 trillion total parameters. Ai2 describes it as a short capacity test rather than a full training run.
The company plans to use an MoE architecture for its next-generation Olmo. Its stated goals include its largest training dataset and longest context window yet.
Sources
- Ai2: Introducing Olmo-core 3: Open, scalable training infrastructure for large MoEs. Source of the benchmarks and the lead chart, published by Ai2.
