Ai2 opens Olmo-core 3 for trillion parameter MoEs

The Allen Institute for AI released Olmo-core 3, a significant upgrade to its open LLM training framework with a redesigned mixture-of-experts training system. The goal is to scale MoE training into the trillion parameter range without losing the efficiency advantage that sparse models are supposed to provide in the first place.

The benchmark that demonstrates the intent is a single configuration. Ai2 increased the expert pool from 8 to 128 while still selecting only four experts per token, keeping active parameters per token roughly fixed at about 3.2B. Total parameter capacity grew from 4.6B to 47B while training throughput fell by less than 5%.

The engineering

The switch is from fully sharded data parallelism, which gathers and reshards weights for each small batch, to a system based on distributed data parallelism that keeps experts resident on GPUs and routes data to them. On a preliminary test with eight NVIDIA B300 GPUs, a 47B parameter MoE processed 52,000 tokens per second per GPU, against 19,400 with the earlier FSDP implementation, roughly 2.7 times the throughput.

Supporting techniques include expert parallelism, pipeline parallelism and a distributed optimizer to spread optimizer state, plus rowwise expert parallelism to write routed data straight into expert input buffers, GPU-resident routing so the CPU does not wait on copied metadata, and grouped GEMM to batch many small expert computations. MXFP8 support gives about 21% higher throughput than BF16 on four B300s, with peak active memory falling from 103 GiB to 95 GiB.

Scaling was benchmarked at 1.2 trillion total parameters with 58.36B active per token across 512 GPUs, peaking at 858 TFLOP/s per GPU. With DeepEP v2 the team reached a 2.38 trillion parameter configuration.

The caveats are stated rather than buried, and they are the useful part of the post. The trillion parameter tests used random routing to measure system performance, not model quality, and the 2.38T figure was a short capacity test rather than a sustained training run. The technical report also names a failure mode worth borrowing the name of: token gerrymandering, where a balance score improves even as the real workload becomes less balanced. Lowering expert learning rates did not help in the model family tested, and overlapping communication with computation on separate streams sometimes slowed training end to end rather than speeding it up.

Ai2 says its next generation Olmo will be an MoE, trained on its largest dataset with its longest context so far. Code is on GitHub, with a technical report and an interactive walkthrough.

(Original source: Ai2 on Hugging Face)