Speed-Read (With Video)

This report describes the MoE training stack in Olmo-core, the measurements that guided its design, and the approaches we tried but did not adopt. This guide summarizes the main results; the table at the end points to sections for readers interested in a particular problem.

Sparse compute leaves substantial costs that scale with total parameter size. An MoE model stores many expert networks but sends each token through only a few of them, so the compute per token stays small while the stored model grows. Weight storage, gradient synchronization, and optimizer updates still depend on total parameter count. The fully sharded data parallel (FSDP) with “full-reshard” configuration used for dense Olmo also gathers and frees weights around every microbatch. In our early baselines, adding experts at fixed active size therefore reduced throughput even though the useful work per token barely changed (Figure 4). Section 1.3 names this capacity tax and separates it from the runtime tax of sparse routing and execution; Section 2 introduces the changes we made to reduce both.

DDP keeps model partitions resident; EP moves selected token rows. The stack replaces the full-reshard FSDP configuration used for dense Olmo with a custom distributed data parallel (DDP) path: each rank keeps its partition of the model in place for the whole accumulation window and synchronizes gradients once per optimizer step, with a distributed optimizer sharding the FP32 main weights and optimizer states (Section 4). Expert parallelism (EP) shards the routed experts across ranks and dispatches only the selected token rows to the rank that stores their expert. Our rowwise intra-node transport writes each token row directly into its destination buffer using GPU-initiated one-sided communication through NVSHMEM, avoiding the host-visible split lists that the block all-to-all path needs. A block-oriented DeepEP v2 backend supports expert groups that span nodes (Section 5). Pipeline parallelism (PP) divides layers when one rank cannot hold them all (Section 6). Section 3 shows how DDP, EP, and PP combine, and Section 7 maps their rank groups onto NVLink and InfiniBand.

GPU-resident metadata removes recurring host waits. The original MoE path copied routing counts to the CPU for both all-to-all split sizes and expert matrix-multiplication group sizes. Keeping this metadata on the device removes those synchronization points from the steady-state loop (Section 8). Device-scheduled grouped general matrix multiplication (grouped GEMM) executes the uneven per-expert shapes produced by routing without padding every expert batch to a common capacity. Section 9 measures where it beats the padded alternative and where it does not.

Routers can hack the balancing loss: lower loss, higher imbalance. Section 10 treats routing as a systems problem: what capacity admits, what the load-balancing loss actually optimizes, and how it can reward a failure we call Token Gerrymandering, in which the loss improves while imbalance grows (an analytic example and one checkpoint-branched run, Figure 42).

MXFP8 reduces compute and data-movement costs. MXFP8, an eight-bit format with shared block scales, can cut the time of large GEMMs, the size of selected saved activations, and dispatch payloads, while one FP32 main weight remains the source of truth (Section 11).

Checkpoints preserve the model independently of its parallel layout. Checkpoints describe every saved chunk by its place in the global FP32 tensor, so a run can resume on a different topology (Section 12).

Recompute lets larger recipes fit in memory. Activation recompute reduces the activations retained for the backward pass, lowering peak high-bandwidth memory (HBM) usage. In our controlled comparison, recomputing every block removes nearly two thirds of activation memory and reduces peak memory by about a quarter, at the cost of about a fifth of throughput. We use it in the larger recipes that would otherwise exceed device memory (Section 13).

The training configurations span 16 to 512 GPUs. Section 14 lists measured operating points, from a 12.9-billion-parameter model on 16 GPUs to a 1.2-trillion-parameter model on 512 GPUs, with a short 2.38-trillion-parameter extension. These configurations differ in batch, parallelism, precision, and recompute, so their rates are not a controlled scaling curve. Table 30 states the batch, topology, precision, and recompute policy needed to interpret each number.

Some optimizations traded one cost for another. Section 15 describes approaches we did not adopt: CPU activation offload, two overlap schemes, a persistent MoE megakernel, a natural-layout MXFP8 weight-gradient prototype, and CUDA Graph replay. The section quantifies their trade-offs between concurrency, kernel efficiency, and memory use; notably, overlapping communication with compute can slow concurrent kernels by more than it hides.

The experiments also lead to training-recipe findings. Critical batch size says when a larger batch still helps but not which learning rate to pair with it, and in our internal tuning runs, square-root scaling needed a small empirical correction (Section 16); in our sweep, scaling expert learning rates by sparsity did not help (Section 17). We train from scratch rather than upcycle because, in the published OLMoE study, a dense model’s head start did not persist over a longer run (Section 18). Several ungated shared experts are mathematically one wider shared expert (Section 20), and the rule that merges routed and shared outputs changes the initialization scale (Section 21).

Beyond pretraining, and comparisons with other systems. Section 22 discusses how the same mechanisms could support supervised fine-tuning, preference training, and reinforcement learning; it is a design discussion rather than a post-training benchmark. Section 23 compares the design against Megatron-Core, DeepSpeed-MoE, and DeepEP in terms of parameter movement, communication, and execution costs. Section 24 summarizes the implications for memory use, throughput, and measurement.