Source-linked AI summary
OLMoE: Open Mixture-of-Experts Language Models
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, Hannaneh Hajishirzi
TL;DR
The paper addresses the limited openness and cost of high-performing MoE language models by introducing a fully open sparse MoE and studying its design and routing. OLMoE-1B-7B uses sparse activation during large-scale pretraining, while an instruct version adds instruction and preference tuning; the models achieve strong performance among similarly active-sized models and exceed several larger baselines.
Problem
Most MoEs release limited information about training data, code, or recipes, while high-performing language models remain costly to build and deploy.
Method
The paper develops a fully open sparse MoE, pretrains it on mixed data, adapts it with instruction and preference tuning, and studies MoE design and routing.
Results
OLMoE-1B-7B achieves state-of-the-art performance among models with similar active parameters, while OLMoE-1B-7B-Instruct exceeds larger models including Llama2-13B-Chat and DeepSeekMoE-16B.
Takeaways & Limitations
The fully open release provides models, data, code, logs, and checkpoints for further research on cost-efficient MoEs and their routing behavior.
Takeaways & Limitations
QK-Norm reduces throughput by almost 10%.
Abstract
from arXiv · showhide
We introduce OLMoE, a fully open, state-of-the-art language model leveraging sparse Mixture-of-Experts (MoE). OLMoE-1B-7B has 7 billion (B) parameters but uses only 1B per input token. We pretrain it on 5 trillion tokens and further adapt it to create OLMoE-1B-7B-Instruct. Our models outperform all available models with similar active parameters, even surpassing larger ones like Llama2-13B-Chat and DeepSeekMoE-16B. We present various experiments on MoE training, analyze routing in our model showing high specialization, and open-source all aspects of our work: model weights, training data, code, and logs.
OLMoE
OLMoE-1B-7B is presented as an open MoE model with publicly available model, data, code, and logs. Figure 1 highlights its performance and openness relative to other open MoE and dense language models.
- OLMoE-1B-7B is the model identified for this release.
- Training data, code, and logs are provided through linked public resources.
- Figure 1 compares performance, cost, and openness across open MoE and dense language models.Model names encode active and total parameter counts for MoEs and total parameters for dense models; MMLU is used as a performance summary.
1 Introduction
OLMoE addresses the cost and accessibility challenges of language models by developing a fully open, sparsely activated MoE. The model combines efficient activation with broad release materials and experiments on MoE design and routing.
- Sparsely activated MoEs aim to improve the performance–cost trade-off by activating only a subset of experts for each input.Dense models activate all parameters for every input, whereas MoEs use several experts in each layer and select a subset.
- Most MoEs provide limited information about training data, code, or recipes, leaving important design questions insufficiently documented.The paper identifies choices involving active versus total parameters, expert structure, shared experts, and routing algorithms.
- OLMoE-1B-7B uses 6.9B total parameters, activates 1.3B per token, and is pretrained on 5.1T tokens.Its inference cost is described as similar to dense models with around 1B parameters, although storing the total parameters requires more GPU memory.
- OLMoE-1B-7B-Instruct exceeds several larger instruct models on common benchmarks after instruction and preference tuning.The cited comparisons include Llama2-13B-Chat, OLMo-7B-Instruct (0724), and DeepSeekMoE-16B.
- Controlled experiments examine MoE design choices and find specialized routing behavior in OLMoE-1B-7B.The analysis reports early routing saturation, rare expert co-activation, and domain and vocabulary specialization.
- The project releases training code, intermediate checkpoints, logs, and training data under open-source licenses.Intermediate checkpoints are released every 5000 steps.
2 Pretraining and Adaptation
OLMoE replaces dense feedforward networks with routed expert modules and trains the model from scratch on a mixed corpus. Adaptation then applies instruction tuning followed by preference tuning.
- Pretraining architecture: OLMoE replaces each dense FFN with an MoE module containing multiple smaller FFN experts, of which only a subset is activated per token.The architecture is decoder-only and consists of transformer layers.
- Pretraining architecture: A learned router maps input logits to experts, applies softmax routing probabilities, and combines the selected expert outputs.Each selected expert processes the input, its output is weighted by its routing probability, and the results are summed.
- Model configuration: OLMoE-1B-7B activates 8 of 64 experts per layer, using dropless token-choice routing and 1.3B active out of 6.9B total parameters.Training uses cross-entropy, load-balancing, and router z-loss terms.
- Pretraining data: The pretraining corpus mixes DCLM and Dolma 1.7 sources, including filtered Common Crawl, StarCoder, Algebraic Stack, arXiv, peS2o, and Wikipedia.The resulting dataset is called OLMoE-MIX.
- Pretraining data: The data pipeline removes documents with at least 32 repeated n-grams and applies additional frequency and repository filters to StarCoder.
- Pretraining data: Pretraining runs for 5.133T tokens, with dataset reshuffling and linear learning-rate decay during the final 100B tokens.
- Adaptation: OLMoE-1B-7B-Instruct is created through instruction tuning followed by preference tuning.The instruction-tuning data adds code and mathematics data for downstream coding and math applications.
3 Results
OLMoE-1B-7B is evaluated during pretraining, after pretraining, and after adaptation, showing strong performance at low active-parameter and compute budgets. Its instruction-tuned variant achieves especially strong results on downstream tasks.
- Evaluation setup: The evaluation procedure covers pretraining, post-pretraining, and post-adaptation stages using the OLMES evaluation standard.The paper compares pretrained and adapted models against larger MoEs and dense language models.
- During pretraining: OLMoE-1B-7B reaches better performance across all tasks with less FLOPs than dense OLMo models during pretraining.It matches or outperforms OLMo-7B while using less than half as many training FLOPs and only 1B active parameters.
- After pretraining: OLMoE-1B-7B is the state of the art in its cost regime after pretraining.It performs best among models using less than 2B active parameters and outperforms some dense 7B models despite requiring ∼6–7× less compute per forward pass.
- After adaptation: SFT improves OLMoE-1B-7B on all measured tasks, including a >10× gain on GSM8k.The authors attribute the GSM8k gain likely to additional math data in instruction tuning.
- After adaptation: OLMoE-1B-7B-Instruct has the highest average among all benchmarked models and scores 84% on AlpacaEval.It outperforms Qwen1.5-3B-14B’s chat version and larger dense models such as Llama2-13B-Chat on AlpacaEval.
4 Experimenting with Alternative Design Choices
The experiments examine MoE design choices including compute efficiency, expert granularity, shared experts, routing, and auxiliary losses. They favor fine-grained dropless token routing while highlighting trade-offs involving throughput, flexibility, and specialization.
- MoE vs. Dense: MoEs reach dense-model performance with ∼3× fewer tokens, but process 23,600 versus 37,500 tokens per second per GPU for dense models.The controlled comparison uses models with the same active parameters and attributes lower throughput to the MoE’s 7B total-parameter memory overhead.
- Shared experts: Removing one routed expert for a shared expert eliminates almost 90% of possible combinations and slightly worsens performance.The authors therefore do not use shared experts, favoring flexibility through more routed-expert combinations.
- Expert Choice vs. Token Choice: Dropless token-based routing outperforms expert-based routing in the authors’ experiments.Token choice assigns a fixed number of experts to each token, whereas expert choice provides perfect load balance but can drop or duplicate tokens and is difficult for autoregressive generation.
- Load Balancing Loss: The load balancing loss encourages equal expert assignment but can constrain specialization by forcing experts to be used approximately equally.The authors use it with weight 0.01 while identifying removal of the loss as future research.
- Router Z-loss: Router z-loss improves stability and quality across training loss, validation loss, and downstream performance despite reducing throughput by ∼2%.OLMoE-1B-7B uses router z-loss with weight 0.001.
MMLU Var (Acc %)
The experiments examine optimization and normalization choices for OLMoE, including their effects on stability, performance, throughput, and adaptation. The final configuration favors RMSNorm, QK-Norm, router z-loss, and weight decay on all parameters despite selected throughput costs.
- Normalization: RMSNorm improves performance over non-parametric layer normalization despite reducing training throughput by 15%.The authors attribute this advantage partly to fewer disruptive gradient spikes and include RMSNorm parameters in weight decay.
- Weight decay: Decaying embedding parameters has only a minor performance impact, with decay slightly better, so OLMoE applies weight decay to all parameters.This includes both embedding and RMSNorm parameters.
- Normalization: QK-Norm improves stability and performance relative to omitting normalization after query and key projections.With RMSNorm, it also slightly improves training loss and prevents a large gradient-norm spike, while reducing throughput by almost 10%.
- Optimization: Reducing AdamW epsilon from 1E-05 to 1E-08 significantly improves performance while maintaining stability.The final run uses eps = 1E-08.
- Adaptation: Instruction-tuning experiments find no conclusive evidence favoring load-balancing or router z-loss auxiliary losses, with only small performance differences.The cited adaptation setup compares auxiliary-loss choices during instruction tuning.
5 MoE Analysis
OLMoE’s routing analysis finds that expert assignments stabilize early and exhibit specialization by domain, vocabulary, and layer, while within-layer expert co-activation is generally weak. These analyses define four properties that characterize how the sparse model uses its experts.
- Router Saturation: Router saturation measures how closely intermediate-checkpoint expert assignments match the final checkpoint assignments for the same tokens.The analysis considers top-k expert sets, with N denoting the token count and k the number of activated experts.
- Router Saturation: After 1% of pretraining, up to approximately 60% of top-8 routing has already saturated, rising to approximately 80% by 40%.Top-1 routing saturates more slowly, and later layers stabilize earlier than layer 0.
- Expert Co-activation: Expert co-activation is generally weak within layers, with only a few exceptions and several small groups in layers 7 and 15.High co-activation can indicate redundancy or motivate placing experts together to reduce communication costs.
- Domain Specialization: Domain routing in OLMoE frequently deviates from random chance, including near-100% specialization for arXiv in one layer-0 expert.Generic C4 data is more balanced, whereas Mixtral shows little domain specialization and remains near uniform routing.
- Vocabulary Specialization: Vocabulary specialization measures how often tokens with a given vocabulary ID are routed to a particular expert within a layer.The analysis compares specialization across layers and experts using token IDs and selected top-k routing.
- Vocabulary Specialization: Vocabulary specialization increases in later layers and increasingly reflects predicted output tokens rather than original input tokens.Expert 27 reaches approximately 90% specialization and activates for many non-alphabetic tokens.
6 Related Work
The paper situates OLMoE within ongoing exploration of sparse MoE architectures and efforts toward fully open language-model research. It emphasizes that meaningful openness includes training data and code alongside model weights.
- Advances in MoEs: Mixture-of-Experts sparsity remains an actively explored architectural modification, while many current language models still use dense architectures.The related-work discussion places MoE alongside widely adopted Transformer changes such as decoder-only training, SwiGLU, RoPE, MQA/GQA, and RMSNorm.
- Open LMs: Fully open language-model releases provide model weights, training data, and code, enabling broader scientific research than weight-only releases.The paper lists several fully open model families and distinguishes them from models released primarily with weights.
- Open LMs: Among MoE models, the paper identifies OpenMoE as a fully open effort but presents OLMoE as a fully open MoE with state-of-the-art performance among similarly sized models.OLMoE releases model weights, data, code, and logs.
7 Conclusion
OLMoE releases models and instruction-tuned models alongside the data, code, logs, and other artifacts needed to support open research and analysis.
- OLMoE-1B-7B and OLMoE-1B-7B-Instruct achieve state-of-the-art performance among models with similar active parameters, also outperforming larger models including DeepSeekMoE-16B and Llama2-13B-Chat.
- The release includes model, data, code, and logs, with Table 9 cataloging the artifacts and linking each named artifact to its retrieval URL.
- The paper analyzes router saturation, expert co-activation, and domain and vocabulary specialization to support further study of MoE behavior.
B Training Configuration
The training configuration combines a highly granular dropless MoE with stability and convergence-oriented choices, followed by supervised and preference-based adaptation. Pretraining used substantial hardware and a long training run.
- Pretraining: OLMoE-1B-7B uses 64 small experts with an FFN dimension of 1,024, plus router z-loss and load-balancing loss.
- Pretraining: The model uses truncated-normal initialization, QK normalization, RMSNorm, and a reduced AdamW epsilon to improve stability or speed convergence.
- Pretraining: OLMoE-1B-7B is pretrained with AdamW, ZeRO through PyTorch FSDP, mixed precision, and a training duration longer than comparable configurations.
- Adaptation: Adaptation uses Open Instruct data, filters SFT samples below 4,096 tokens, aggregates loss at token level, and trains for two epochs.
- Hardware: Pretraining uses 256 H100 GPUs for approximately 10 days, while instruction and preference tuning use 32 H100 GPUs.
C Evaluation Setup
The evaluation setup measures pretrained and adapted models across broad task suites using documented, reproducible protocols, while comparing OLMoE with related models and released artifacts.
- During pretraining: During pretraining, evaluation adds CommonsenseQA, PIQA, and multiple MMLU implementations to the in-loop task suite.
- After pretraining: After pretraining, OLMES provides an open, practical, documented standard for reproducible language-model evaluation across smaller base models.
- After pretraining: DCLM results follow the authors’ released evaluation code and distinguish Core low-variance tasks from Extended heavy tasks.
- After adaptation: Post-adaptation evaluation covers coding, reasoning, world knowledge, instruction following, precise instruction following, and safety.
- Openness comparisons: The paper compares model openness across model, data, code, and logs, including licensing and whether training mixtures or implementations are released.
- Reported evaluations: Figures and tables document pretraining losses, evaluation against current OLMo models, OLMES results, and DCLM metrics.
F Additional Experiments
Additional experiments test data mixtures, numerical precision, upcycling noise, shared layers, adaptation choices, and routing specialization. These studies identify several design trade-offs and characterize expert behavior across domains and layers.
- Data and training choices: Adding Reddit or FLAN to the pretraining mixture does not produce consistent gains, so neither is included in the final mix.
- Data and training choices: Computing the load-balancing loss in FP32 does not reduce stability spikes, so the model retains BF16 computation.
- Data and training choices: After 700 billion tokens, no-noise upcycling performs slightly better, although both variants appear to converge to the same performance.
- Shared layers: Layer-shared MoEs perform similarly to regular dense models, retain a small advantage on validation loss and HellaSwag, and reduce training throughput by around 20%.
- Adaptation: KTO experiments select the RMS optimizer and the 5,000-step checkpoint for adaptation.
- Routing analysis: Routing analyses visualize vocabulary and domain specialization, load imbalance after adaptation, and token-routing flows across layers for OLMoE and Mixtral.
H Limitations and Future Work
The release identifies limits in active parameter capacity, training data scale, modality, and language coverage, while describing a subsequent OLMoE-1B-7B-0125 update with improved evaluation performance and revised adaptation.
- Limitations: 1B active parameters constrain OLMoE-1B-7B's capabilities relative to models using more than seven times as many parameters.The paper notes that adding parameters could improve performance by allowing more than 1B parameters per input.
- Limitations: 5 trillion training tokens may be insufficiently explored for MoEs, since some dense models train on 15 trillion tokens and longer MoE training remains understudied.The paper reports a token multiplier of around 5,000 and says the effectiveness of further overtraining for MoEs requires more research.
- Limitations: OLMoE-1B-7B is text-only, so it cannot accept or produce image or audio inputs and outputs.The paper identifies this as limiting utility for multimodal use cases.
- Limitations: The predominantly English pretraining corpus and exclusively English evaluation limit the model's usefulness for non-English language-model research.The paper suggests future data additions may include more non-English content.
- Future Work: +1.6 average performance and +2.1 on MMLU mark OLMoE-1B-7B-0125 as an improvement over the previous iteration.The update uses a curated DOLMINO annealing mix and an updated post-training recipe featuring improved SFT, DPO sampling, and PPO with verifiers.
- Future Work: The updated models and datasets are freely available on the Hugging Face hub.The release also includes references to the announcement and corrected or expanded materials.