AI2 Releases Olmo-core 3, an Open Training Framework for Trillion-Parameter MoE Models

AI2 has released Olmo-core 3, a major upgrade to its framework for developing large language models, featuring a redesigned open mixture-of-experts (MoE) training system. The framework is designed to scale MoE training into the trillion-parameter range while preserving computational efficiency, and it underpins the next generation of Olmo models.

MoE models offer a more efficient training approach by activating only a subset of parameters per token, but communication and coordination overhead can erode their computational advantage as they scale. In one benchmark described by AI2, the expert pool was increased from 8 to 128 while still selecting only four experts per token. Total parameter capacity grew from 4.6B to 47B, while training throughput fell by less than 5%. The same infrastructure has also been benchmarked at over one trillion total parameters.

The new system switches from fully sharded data parallelism (FSDP) to a distributed data parallelism (DDP)-based design, keeping experts resident on GPUs and routing data to them to avoid repeated weight gathering. It combines expert, pipeline, and distributed optimizer parallelism with optimizations such as rowwise expert parallelism, GPU-resident routing, and grouped GEMM. Olmo-core 3 also supports MXFP8, a lower-precision format that improved end-to-end training throughput by about 21% over BF16 in a controlled benchmark on four NVIDIA B300 GPUs, while reducing peak active memory from 103 GiB to 95 GiB.

In tests on NVIDIA B300 GPUs, AI2 benchmarked a 1.2-trillion-parameter model with 58.36 billion active parameters per token across 512 GPUs, reaching a peak throughput of 858 TFLOP/s/GPU. The team also experimented with DeepEP v2 in a short-capacity test reaching 2.38 trillion total parameters. Olmo-core 3 is fully open, with the technical report and code available on GitHub.