Allen Institute for AI Releases Olmo-core 3 for Trillion-Parameter MoE Model Training

Allen Institute for AI Releases Olmo-core 3 for Trillion-Parameter MoE Model Training The Allen Institute for AI has launched Olmo-core 3, an open-source training infrastructure redesign engineered to scale Mixture-of-Experts (MoE) architectures into the trillion-parameter regime while preserving throughput and resource efficiency. Replaces an earlier FSDP-based architecture with distributed data parallelism (DDP) that keeps experts GPU-resident, achieving a 2.7x throughput increase in a 47-billion-parameter test. Maintains high throughput across capacity scaling, dropping less than 5% in throughput when expanding an expert pool from 8 to 128 (4.6B to 47B total parameters). Integrates expert and pipeline parallelism, distributed optimizers, GPU-resident routing, grouped GEMM, and optional MXFP8 low-precision execution. Demonstrates scaling to a 1.2-trillion-parameter model reaching 858 TFLOP/s/GPU across 512 NVIDIA B300 GPUs, alongside a capacity test reaching 2.38 trillion parameters using DeepEP v2. Identifies training behaviors including 'token gerrymandering' and cases where overlapping communication and computation harmed end-to-end throughput. Architectural Transition from FSDP to DDP Prior MoE training implementations in Olmo-core relied on fully sharded data parallelism (FSDP), which gathered and resharded expert weights for every small batch of data. Olmo-core 3 shifts to a distributed data parallel (DDP) foundation that keeps model experts resident directly on GPUs and routes data tokens to them instead. In an initial test using eight NVIDIA B300 GPUs on a 47-billion-parameter MoE, this architectural change increased throughput from 19,400 tokens per second per GPU to 52,000 tokens per second per GPU, representing an approximate 2.7-fold increase. Parallelism and Execution Optimizations The framework manages memory and coordination across clusters through a combination of expert parallelism, pipeline parallelism, and a distributed optimizer. To reduce operational overhead, rowwise expert parallelism deposits routed tokens directly into expert input buffers, while GPU-resident routing keeps routing metadata entirely on the GPU to prevent CPU stalls. It also leverages grouped GEMM to batch individual expert matrix computations, alongside support for the MXFP8 format. Benchmarked on four NVIDIA B300 GPUs, MXFP8 improved end-to-end training throughput by roughly 21% compared to BF16 and reduced peak active memory from 103 GiB to 95 GiB, predominantly through efficiencies in feed-forward networks and inter-expert data transfers. Trillion-Parameter Scaling Benchmarks System stress tests evaluated Olmo-core 3 across extensive hardware clusters. In a benchmark configuring 512 NVIDIA B300 GPUs for a 1.2-trillion-parameter model with 58.36 billion active parameters per token, the stack achieved up to 858 TFLOP/s/GPU using random routing. An experimental short-capacity evaluation incorporating DeepEP v2 for cross-GPU expert communications demonstrated the system's ability to host configurations sizing up to 2.38 trillion total parameters. Empirical Training Observations and Future Plans System experiments uncovered multiple implementation trade-offs documented in the project's technical report. The authors identified 'token gerrymandering,' wherein load-balancing metrics improved even as the underlying expert workload grew increasingly unbalanced. They also observed that reducing learning rates for experts seeing fewer tokens offered no benefit, GPU execution times varied depending on input tensor values despite identical dimensions, and overlapping computation with communication across GPU streams sometimes decreased end-to-end throughput. Olmo-core 3 is released as an open framework and serves as the infrastructure for the upcoming MoE-based generation of the Olmo model.

Allen AI推出Olmo-core 3：面向万亿参数MoE模型的开源训练基础设施 Allen AI正式发布Olmo-core 3，重构了开源混合专家（MoE）训练系统，旨在将MoE训练扩展至万亿参数规模并保持计算效率，作为下一代Olmo模型的基础底座全面开源。 Olmo-core 3将架构从早期基于FSDP的反复权重聚合转向基于DDP的方案，使专家驻留于GPU并通过数据路由减少通信开销，初步测试吞吐量提升约2.7倍。 整合了专家并行、流水线并行、分布式优化器、行优先专家并行、GPU驻留路由、分组GEMM以及MXFP8低精度格式等技术。 在512块NVIDIA B300 GPU上完成了1.2万亿参数配置的基准测试，最高吞吐量达858 TFLOP/s/GPU，并在容量测试中触及2.38万亿参数规模。 技术报告揭示了负载平衡评分失真（Token Gerrymandering）、降低专家学习率未见成效以及重叠通信与计算可能反致变慢等实验发现。 重构MoE训练底层架构 相较于此前在Olmo-core中使用FSDP反复收集与重分片模型权重的做法，Olmo-core 3转向基于分布式数据并行（DDP）的系统，让专家模型常驻在GPU显存中并将对应数据路由至专家。在8块NVIDIA B300 GPU测试中，470亿参数MoE模型的单卡吞吐量从旧实现的每秒19,400 token提升至52,000 token，增幅达2.7倍；当专家池从8个扩展到128个（每token激活4个专家）时，模型容量从46亿扩充至470亿参数，而吞吐量下降不足5%。 并行策略与系统级优化 该框架结合了专家并行、流水线并行与分布式优化器来分摊显存占用。同时，行优先专家并行将路由数据直接写入专家输入缓冲区，GPU驻留路由避免了CPU等待元数据回传，分组GEMM合并了离散的小专家计算。此外，Olmo-core 3支持MXFP8低精度格式；在4块B300 GPU的测试中，开启MXFP8后整体训练吞吐量比BF16基准提高约21%，活跃显存峰值从103 GiB降至95 GiB。 万亿参数扩展基准与实验洞察 在512块NVIDIA B300 GPU环境下，Olmo-core 3对1.2万亿总参数（每token激活583.6亿参数）的模型进行了随机路由基准测试，测得最高计算吞吐量达858 TFLOP/s/GPU；结合DeepEP v2通信方案的短期容量测试达到了2.38万亿参数规模。技术报告还指出了实际训练中的系统权衡，包括平衡路由评分改善但实际工作负载更不平衡的“token gerrymandering”现象、降低低频处理token专家的学习率未见提升，以及在多流上重叠通信与计算在某些测试中反而拖慢端到端执行速度等问题。

📄 Tech Report | 💻 Code | 🧩 Interactive demo Today we’re releasing Olmo-core 3 , a significant upgrade to our framework for developing large language models featuring a redesigned open mixture-of-experts (MoE) training system. Olmo-core 3 is designed to scale MoE training into the trillion-parameter range while preserving computational efficiency. It’s one of the core systems behind the next generation of Olmo, and part of our ongoing commitment to open up the tools and training infrastructure behind each new model. Training large AI models takes a lot of compute, driving up costs and energy use and putting advanced model development out of reach for many academic researchers and smaller labs. MoE models offer a more efficient approach—they can contain many more learned components, or parameters, without requiring every input to use all of them. But the full model still has to be stored across GPU memory and updated during training, and directing inputs to the right experts – the specialized components within an MoE – across a cluster creates its own communication and coordination costs. As MoEs grow, those costs can erode much of the computational advantage of using only part of the model for each input. Olmo-core 3 is built to close that gap. In one benchmark, we increased the expert pool from 8 to 128 while still selecting only four experts per token – the small units of text a language model processes – keeping the number of 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 same infrastructure has been benchmarked at over one trillion total parameters. Building a training stack around how MoEs actually work Olmo-core has evolved with each generation of Olmo. Our work on sparse models goes back to OlmoE , which used an MoE architecture with 64 routed experts. Olmo 3 , by contrast, used a dense architecture, meaning nearly all of the model was active for every token and its training stack was built around that design. Olmo-core 3 extends the framework with a training system designed for much larger MoE models. Our earlier MoE implementation in Olmo-core used fully sharded data parallelism (FSDP), configured to gather and reshard model weights for each small batch of training data. Olmo-core 3 switches to a system based on distributed data parallelism (DDP) . It keeps experts resident on GPUs and routes the relevant data to them, avoiding that repeated weight gathering. NVIDIA’s Megatron-Core is an established option for training large MoEs. Olmo-core 3 brings an integrated MoE training stack to the framework behind Olmo, with a redesign that improves throughput over our earlier FSDP-based implementation. In a preliminary test on eight NVIDIA B300 GPUs, a 47-billion-parameter MoE processed 52,000 tokens per second per GPU with the new stack, compared with 19,400 using our earlier implementation—about 2.7× the throughput. Scaling and optimizing MoE training Olmo-core 3 combines several techniques for distributing large MoEs across GPU clusters with optimizations that make routing and computation more efficient. Three techniques determine how the model and its training state are split across hardware: Expert parallelism spreads the experts across GPUs, so each GPU stores only part of the full expert pool. Pipeline parallelism splits the model’s layers – the successive stages that transform an input – across groups of GPUs, reducing how much of the model each GPU needs to keep in memory. A distributed optimizer spreads the optimizer state – the additional data used to calculate and apply updates during training – across GPUs instead of storing a full copy on every GPU. Together, these techniques allow an MoE to scale without requiring every GPU to keep the entire model and its training state in memory. Olmo-core 3 also reduces the cost of routing data to the right experts and running their computations. Rowwise expert parallelism places routed data directly into expert input buffers, minimizing the extra work needed to rearrange it. GPU-resident routing keeps routing metadata on the GPUs, so the CPU can queue work without waiting for that information to be copied back. And grouped GEMM combines many small expert computations so GPUs can execute them more efficiently. Finally, Olmo-core 3 supports MXFP8 , a lower-precision number format that represents some values with fewer bits. This can reduce computation and the amount of data moved between GPUs, as long as those savings outweigh the cost of converting between number formats. We measured MXFP8’s effect on end-to-end training throughput in a controlled benchmark on four NVIDIA B300 GPUs, with work distributed uniformly across experts. With MXFP8 enabled across the parts of the system where it helped most, training throughput was about 21% higher than with BF16, the higher-precision format we used as our baseline, while peak active memory fell from 103 GiB to 95 GiB. Most of the gain came from feed-forward computation and moving data between experts rather than attention alone. These techniques and optimizations have to work together. Speeding up one part of training can create costs elsewhere; faster computation may require more data movement, while moving fewer bits may not help if converting the data takes too long. Olmo-core 3 is built around those trade-offs across the full training process, giving us – and researchers using the open stack – control over how the pieces fit together. Explore our interactive walkthrough to see how data, expert, and pipeline parallelism work together to scale MoE training—from a single GPU to many. Scaling into the trillion-parameter range We’ve benchmarked Olmo-core 3 across a range of configurations on NVIDIA B300 GPUs, including a 1.2-trillion-parameter model with 58.36 billion parameters active per token across 512 GPUs. Its highest observed throughput was 858 TFLOP/s/GPU—a measure of useful model computation per second on each GPU. These tests used random routing to measure system performance, rather than the quality of a trained model. We’ve also experimented with DeepEP v2, an alternative way of handling communication between experts across GPUs, reaching a configuration with 2.38 trillion total parameters. This was a short-capacity test rather than a full training run, so it demonstrates the scale Olmo-core 3 can reach rather than sustained training performance. At these scales, systems performance is only part of the picture. Our technical report also documents experiments that informed how we train MoEs and measure their performance. For example: A score intended to encourage balanced routing could improve even as the actual workload became less balanced. We call this failure token gerrymandering . Lowering experts’ learning rates – the size of their training updates – because they process fewer tokens did not improve results in the model family we tested. GPU calculations took different amounts of time when the values being processed changed, even with the same matrix dimensions. Performance comparisons therefore need matching input values as well as matching shapes. Overlapping communication and computation on separate GPU streams did not always make training faster. In some tests, it slowed end-to-end execution—a reminder that more overlap does not necessarily mean higher throughput. The report explains these findings alongside the approaches we tested and chose not to adopt. Built for the next generation of Olmo, open for everyone Olmo-core 3 is the foundation for what we’re building next. Our next-generation Olmo will use an MoE architecture, and we’re aiming for it to be our most capable Olmo yet, trained on our largest dataset and with our longest context window. The new stack lets us scale beyond our previous MoE work while giving us more flexibility to adapt training as models and hardware evolve. And it’s fully open—researchers and developers can use Olmo-core 3 to train their own MoEs, adapt it to different hardware, and experiment with routing, parallelism, and other parts of the system. That’s part of how we think about open model development—model weights are more useful when the infrastructure and training decisions behind them are open too. For a deeper look at the systems design, experiments, ablations, and approaches we tested along the way, read our technical report and explore Olmo-core 3 on GitHub .