Abstract
Mixture-of-Experts (MoE) models offer more total parameters without proportionally increasing per-token compute: each token activates only a few experts. However, training cost depends not only on active compute but also on many other factors, including bytes moved, memory held, and kernels launched; many of those costs follow total parameter size rather than active compute. On infrastructure built for dense models, this mismatch is expensive: in our early baselines, throughput fell as stored expert parameter size grew, even though the useful work per token stayed nearly fixed.
To close this gap, we present a redesigned training stack in Olmo-core, which is optimized for efficient and scalable MoE training. We switch from the Fully Sharded Data Parallel (FSDP) used by previous Olmo-core versions to a parallelization architecture built on Distributed Data Parallel (DDP). On top of it, we add Expert Parallelism (EP), Pipeline Parallelism (PP), a distributed optimizer, and activation recompute to support large-scale models.
We bring together a suite of optimizations to make MoE training more efficient. For MoE layers, we propose NVSHMEM-based rowwise EP, which writes each token directly into its expert’s buffer, avoiding permutation overhead in dispatch and combine operations. We make EP fully synchronization-free by keeping routing metadata on the GPU and using device-scheduled grouped GEMMs. Closely integrated MXFP8 raises throughput by cutting compute, activation-memory, and transport costs, and topology-agnostic checkpoints let the parallel layout change between runs.
On NVL8 B300 nodes, the stack supports training configurations from 12.9 billion to 1.2 trillion total parameters on up to 512 GPUs, reaching 858 useful-model TFLOP/s/GPU at 1.2 trillion with MXFP8 and per-layer recompute; an experimental capacity test with the optional DeepEP v2 backend reaches 2.38 trillion.
We also explain the reasoning behind each design choice and share findings that can inform other training systems. For example, we identify a failure pattern we call Token Gerrymandering, in which the MoE router can learn to hack the load-balancing loss, and find that overlapping communication with computation can slow the overlapped kernels by more than it hides. These are two of many findings in the report. Together, the stack and these findings offer practical insights into building efficient, scalable MoE training systems.