OLMoE: Open Mixture-of-Experts Language Models

## Abstract

Abstract 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. Weights https://hf.co/allenai/OLMoE-1B-7B-0924 Data {https://hf.co/datasets/allenai/OLMoE-mix-0924} Code https://github.com/allenai/OLMoE Logs https://wandb.ai/ai2-llm/olmoe/reports/ OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3

## 1 Introduction

Despite significant advances in Large Language Models (LMs) on various tasks, there remains a clear trade-off between performance and cost in both training and inference. High-performing LMs are inaccessible for many academics and open-source developers as they are prohibitively expensive to build and deploy.(^1^11For example, even with 16 H100 GPUs and several optimizations, Llama 3 405B only achieves a decoding throughput of around 100 tokens per second [].) One approach to improve the cost-performance trade-off lies in using sparsely-activated Mixture-of-Experts (MoEs) []. MoEs have several experts in each layer, only a subset of which is activated at a time (see Figure 2). This makes MoEs significantly more efficient than dense models with a similar number of total parameters, which activate all parameters for every input []. For this reason, industry frontier models use MoEs including Gemini-1.5 [] and reportedly GPT-4 [].

Most MoE models, however, are closed-source: while some have publicly released model weights [, , , , ], they offer limited to no information about their training data, code, or recipes (see Figure 1). While there have been prior efforts in making language modeling research fully accessible [, , , , , ], they have been largely limited to dense LMs. This comes despite MoEs requiring more openness as they add complex new design questions to LMs, such as how many total versus active parameters to use, whether to use many small or few large experts, if experts should be shared, and what routing algorithm to use. The lack of open resources and findings about these details prevents the field from building cost-efficient open MoEs that approach the capabilities of closed-source frontier models.

To address these issues, we introduce OLMoE, a fully open Mixture-of-Experts language model with state-of-the-art performance among similarly-sized models. In particular, we pretrain OLMoE-1B-7B for 5.1 trillion tokens with 6.9B total parameters, of which only 1.3B are activated for each input token. This leads to a similar inference cost as using dense models with around 1B parameters, such as OLMo 1B [] or TinyLlama 1B [], but requires more GPU memory to store its 7B total parameters. Our experiments show that MoEs train 2 faster than dense LMs with equivalent active parameters. In Figure 1, we show that OLMoE-1B-7B significantly outperforms all open 1B models and displays competitive performance to dense models with significantly higher inference costs and memory storage (e.g., similar MMLU scores to Llama2-13B, which is 10 more costly). Via instruction- and preference tuning, we create OLMoE-1B-7B-Instruct, which we find exceeds various larger instruct models including Llama2-13B-Chat [], OLMo-7B-Instruct (0724), and DeepSeekMoE-16B [] on common benchmarks (MMLU, GSM8k, HumanEval, etc.).

Our comprehensive set of controlled experiments highlights key design choices for MoEs (see Table 1) and LMs in general. One critical design decision for making MoEs performant is the use of fine-grained routing with granular experts []: we employ 64 small experts in each layer with 8 being activated. The choice of routing algorithm is also important: we find dropless [] token-based routing [] outperforms expert-based routing []. Our findings also include those that challenge prior work, such as the ineffectiveness of shared experts [] and the limited benefits of sparsely upcycling a pretrained dense LM into an MoE [] unless under small compute budgets. Finally, we analyze the routing behavior in OLMoE-1B-7B, finding that routing saturates early in pretraining, experts are rarely co-activated, and experts exhibit domain and vocabulary specialization.

We hope our fully open MoE facilitates more research and analysis to improve our understanding of these models. We release training code, intermediate checkpoints (every 5000 steps), training logs, and training data under open-source licenses (Apache 2.0 [http://www.apache.org/licenses/LICENSE-2.0](http://www.apache.org/licenses/LICENSE-2.0) or ODC-By 1.0 [https://opendatacommons.org/licenses/by/1-0/](https://opendatacommons.org/licenses/by/1-0/)).

## 2 Pretraining and Adaptation

Figure: Figure 2: Comparison of the architecture of dense LMs and MoE models like OLMoE. The figure excludes some details, e.g., OLMoE-1B-7B also uses QK-Norm (§4.2.5).
Refer to caption: x4.png

**Table 1: Key MoE design choices and our setup for OLMoE-1B-7B based on our experiments. Full configuration for OLMoE-1B-7B is in Appendix B.**
| Design choice | Description | Exper-iment | OLMoE-1B-7B |
| --- | --- | --- | --- |
| Active params | # active parameters per input token | §4.1.1 | 1.3B active |
| Total params | Total # of parameters in the model | §4.1.1 | 6.9B total |
| Expert granularity | Using fine-grained small experts vs. a few large experts [] | §4.1.2 | 64 small experts with 8 activated |
| Expert sharing | Whether or not to include a shared expert [] | §4.1.3 | No shared expert |
| Routing algorithm | How inputs are assigned to experts, e.g., assignment on a per token basis (e.g., 2 experts per token) or per expert basis (e.g., 2 tokens per expert), and whether or not all tokens get assigned or some get dropped [, ] | §4.1.4 | Dropless [] MoE with token choice |
| Sparse upcycling | Whether to start from a dense model [, ] | §4.1.5 | Not used |
| Load balancing loss | Auxiliary loss to penalize unequal assignment to experts that may harm performance [] | §4.1.6 | Used with weight 0.01 |
| Router z-loss | Auxiliary loss to penalize large logits in the router that may cause instabilities [] | §4.1.7 | Used with weight 0.001 |

**Table 2: Composition of the pretraining data for OLMoE-1B-7B. StarCoder, peS2o, and Wikipedia parts come from Dolma 1.7 []. Links to our data are in Appendix A.**
| Source | Doc Type | GPT-NeoX | Words<br>(billions) | UTF-8 | Documents<br>(millions) |
| --- | --- | --- | --- | --- | --- |
| tokens | bytes |  |  |  |  |
| (billions) | (GB) |  |  |  |  |
| DCLM-Baseline [] | web pages | 3,860 | 3,380 | 16,700 | 2,950 |
| StarCoder [, ] | code | 101 | 63.9 | 325 | 78.7 |
| peS2o [, ] | STEM papers | 57.2 | 51.3 | 268 | 38.8 |
| arXiv [] | STEM papers | 21.1 | 23.5 | 88.8 | 1.55 |
| OpenWebMath [] | math web pages | 12.7 | 10.2 | 42.4 | 2.91 |
| Algebraic Stack [] | math proofs code | 12.6 | 9.6 | 39.3 | 2.83 |
| English Wikipedia<br>& Wikibooks [] | encyclopedic | 3.69 | 3.16 | 16.2 | 6.17 |
| Total | 4,060 | 3,530 | 17,400 | 3,080 |  |

##### Pretraining architecture

OLMoE is a decoder-only LM consisting of $N_{L}$ transformer [] layers. The feedforward network (FFN) in dense models like OLMo [], is replaced with an MoE module consisting of $N_{E}$ smaller FFN modules called experts, of which a subset of $k$ experts are activated for each processed input token $x$ (also see Figure 2):

$$ $\text{MoE module}(x)=\tsum\slimits@_{i=1}^{\text{Top}-k(r(x))}\mathrm{softmax}\left(r(x)\right)_{i}{E_{i}(x)}$ (1) $$

where $r$, called the router, is a learned linear layer mapping from the input logits to the chosen $k$ experts. A softmax is applied to the router outputs to compute routing probabilities for all $N_{E}$ experts. Each selected expert $E_{i}$ processes the input $x$, the output of which is then multiplied with its respective routing probability. The results are then summed across all chosen Top-$k$ experts to constitute the output of the MoE module for a single layer of the model out of its $N_{L}$ total layers. Key decisions in designing an MoE model include determining the number of activated and total parameters, the design of the experts (e.g., granularity, whether or not to include shared experts), and the choice of the routing algorithm. Moreover, training an MoE model can involve initializing from a dense model (sparse upcycling) and changing the training objective, such as including auxiliary load balancing and router z-losses. Experiments related to these design choices are in §4.1; Table 1 shows our final decisions.

In summary, we use 1.3B active parameters out of a total of 6.9B, with 8 activated experts out of 64 per layer. We use dropless token choice routing []: for each input token, the learned router network determines 8 experts to process it. We train OLMoE-1B-7B from scratch with two auxiliary losses: load balancing loss ($\mathcal{L}_{LB}$) [] and router z-loss ($\mathcal{L}_{RZ}$) [], which we define and experiment with in §4.1.6 and §4.1.7, respectively. We multiply them with respective loss weights, $\alpha$ and $\beta$, and sum them linearly with the cross entropy loss ($\mathcal{L}_{\textit{CE}}$) to arrive at our final training loss:

$$ $\mathcal{L}=\mathcal{L}_{\textit{CE}}+\alpha\mathcal{L}_{\textit{LB}}+\beta\mathcal{L}_{\textit{RZ}}$ (2) $$

Our full pretraining configuration for OLMoE-1B-7B is in Appendix B.

**Table 3: Adaptation training data for OLMoE-1B-7B. Links to our data are in Appendix A.**
| Source | Domain | Samples |
| --- | --- | --- |
| Instruction Tuning |  |  |
| Tulu 2 SFT Mix [] | Various | 326,154 |
| No Robots [] | Various | 9,500 |
| CodeFeedback-Filtered-Instruction [] | Coding | 156,526 |
| MetaMathQA [] | Math | 98,750 |
| Advanced (non-chat) subset of Daring Anteater [] | Various | 17,082 |
| Preference Tuning (DPO []) |  |  |
| UltraFeedback [] binarized and filtered for TruthfulQA [] contamination | Various | 60,800 |

##### Pretraining data

We use a mix of data from DCLM [] and Dolma 1.7 [], which includes the following: (1) a quality-filtered subset of Common Crawl, referred to as DCLM-Baseline, (2) StarCoder, Algebraic Stack and arXiv, used in both DCLM and Dolma 1.7, and (3) peS2o and Wikipedia from Dolma 1.7. We refer to our pretraining dataset as OLMoE-Mix.

To all sources above, we apply a filter that removes all documents with a sequence of 32 or more repeated n-grams, where an n-gram is any span of 1 to 13 tokens.
For the StarCoder subset, we also remove any document that is either from a repository with fewer than 2 stars on GitHub, or whose most frequent word constitutes over 30% of the document, or whose top-2 most frequent words constitute over 50% of the document.

We shuffle all samples randomly at the beginning of each epoch and train for a total of 5.133T tokens (1.3 epochs following Muennighoff et al. []). During our annealing phase (final 100B tokens) we first reshuffle the entire dataset and then linearly decay the learning rate to 0, following prior work [, ]. Our pretraining data statistics are in Table 2.

##### Adaptation

We create OLMoE-1B-7B-Instruct by following a standard adaptation recipe split into instruction tuning [, , , , ] followed by preference tuning [, , , ] building on prior open models [, , ]. In our instruction tuning dataset, we add more code and math data to boost performance on downstream coding and math applications. Other models, such as GPT-4 [] and Llama 3 [], instead include samples from math datasets like GSM8k [] or MATH [] during pretraining. We also include No Robots and a subset of Daring Anteater as they are of high quality and add diversity, two key factors for successful adaptation [, , , ]. We describe our adaptation datasets in Table 3 and hyperparameters in Appendix B.

## 3 Results

Our evaluation procedure consists of three parts: During pretraining, After pretraining, and After adaptation. We detail the setup for each in Appendix C.

Figure: Figure 3: Evaluation of OLMoE-1B-7B and the current best OLMo models during pretraining. OLMoE-1B-7B differs from the OLMo models in its MoE architecture, several training hyperparameters, and its training dataset, see §2. A version of this plot with tokens as the x-axis and markers where annealing starts is in Appendix E. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz](https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz)
Refer to caption: x5.png

##### During pretraining

In Figure 3 we benchmark the performance of OLMoE-1B-7B during pretraining with the current best OLMo models [] on commonly used downstream tasks. We find that across all tasks OLMoE-1B-7B reaches better performance with less compute (FLOPs) than the dense OLMo models. OLMoE-1B-7B matches or outperforms OLMo-7B at the end of training despite OLMoE-1B-7B having used less than half as many FLOPs for training and using only 1B active parameters. This is likely a result of the dataset and modeling changes we make to the OLMo setup including MoE-related changes, stability, and performance improvements, outlined in Appendix B. Appendix E contains training and validation loss plots showing very smooth loss curves without major loss spikes during the 5T tokens of our pretraining.

**Table 4: OLMoE-1B-7B after pretraining versus larger MoEs and dense LMs. We compare with dense LMs that assimilate OLMoE-1B-7B either in active parameters (1B, approximates speed and cost) or total parameters (7B, approximates memory requirements). Model names contain rounded parameter counts: model-active-total for MoEs and model-total for dense LMs (this leads to some differences to official names, e.g., while called “Gemma2-2B” it actually has 2.6B active and total parameters []). Chall. = Challenge. We run all evaluations ourselves with 5 few-shots, see Appendix C for details.**
|  | Active | Open | MMLU | Hella- | ARC- | ARC- | PIQA | Wino- |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | params | Data | Swag | Chall. | Easy | Grande |  |  |
| LMs with 7-9B active parameters |  |  |  |  |  |  |  |  |
| Llama2-7B [] | 6.7B |  | 46.2 | 78.9 | 54.2 | 84.0 | 77.5 | 71.7 |
| OLMo-7B (0724) [] | 6.9B |  | 54.9 | 80.5 | 68.0 | 85.7 | 79.3 | 73.2 |
| Mistral-7B [] | 7.3B |  | 64.0 | 83.0 | 78.6 | 90.8 | 82.8 | 77.9 |
| DCLM-7B [] | 6.9B |  | 64.4 | 82.3 | 79.8 | 92.3 | 80.1 | 77.3 |
| Llama3.1-8B [] | 8.0B |  | 66.9 | 81.6 | 79.5 | 91.7 | 81.1 | 76.6 |
| Gemma2-9B [] | 9.2B |  | 70.6 | 87.3 | 89.5 | 95.5 | 86.1 | 78.8 |
| LMs with 2-3B active parameters |  |  |  |  |  |  |  |  |
| OpenMoE-3B-9B [] | 2.9B |  | 27.4 | 44.4 | 29.3 | 50.6 | 63.3 | 51.9 |
| StableLM-2B [] | 1.6B |  | 40.4 | 70.3 | 50.6 | 75.3 | 75.6 | 65.8 |
| DeepSeek-3B-16B [] | 2.9B |  | 45.5 | 80.4 | 53.4 | 82.7 | 80.1 | 73.2 |
| JetMoE-2B-9B [] | 2.2B |  | 49.1 | 81.7 | 61.4 | 81.9 | 80.3 | 70.7 |
| Gemma2-3B [] | 2.6B |  | 53.3 | 74.6 | 67.5 | 84.3 | 78.5 | 71.8 |
| Qwen1.5-3B-14B [] | 2.7B |  | 62.4 | 80.0 | 77.4 | 91.6 | 81.0 | 72.3 |
| LMs with 1B active parameters |  |  |  |  |  |  |  |  |
| Pythia-1B [] | 1.1B |  | 31.1 | 48.0 | 31.4 | 63.4 | 68.9 | 52.7 |
| OLMo-1B (0724) [] | 1.3B |  | 32.1 | 67.5 | 36.4 | 53.5 | 74.0 | 62.9 |
| TinyLlama-1B [] | 1.1B |  | 33.6 | 60.8 | 38.1 | 69.5 | 71.7 | 60.1 |
| DCLM-1B [] | 1.4B |  | 48.5 | 75.1 | 57.6 | 79.5 | 76.6 | 68.1 |
| OLMoE-1B-7B | 1.3B |  | 54.1 | 80.0 | 62.1 | 84.2 | 79.8 | 70.2 |

**Table 5: OLMoE-1B-7B after adaptation versus other models. We find the JetMoE chat model ([https://hf.co/jetmoe/jetmoe-8b-chat](https://hf.co/jetmoe/jetmoe-8b-chat)) is broken leading to random scores thus we exclude it. Names contain rounded parameter counts: model-active-total for MoEs and model-total for dense LMs. We run all evaluations ourselves (Appendix C). Models use different mixes for adaptation, e.g., OLMoE is trained on an improved version of the pipeline used for OLMo models.**
|  |  |  |  | Human- | Alpaca- |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Task () | MMLU | GSM8k | BBH | Eval | Eval 1.0 | XSTest | IFEval | Avg |
| Setup () | 0-shot | 8-shot CoT | 3-shot | 0-shot | 0-shot | 0-shot | 0-shot |  |
| Metric () | EM | EM | EM | Pass@10 | %win | F1 | Loose Acc |  |
| OLMo-1B (0724) | 25.0 | 7.0 | 22.5 | 16.0 | - | 67.6 | 20.5 | - |
| +SFT | 36.0 | 12.5 | 27.2 | 21.2 | 41.5 | 81.9 | 26.1 | 35.9 |
| +DPO | 36.7 | 12.5 | 30.6 | 22.0 | 50.9 | 79.8 | 24.2 | 37.4 |
| OLMo-7B (0724) | 50.8 | 32.5 | 36.9 | 32.3 | - | 80.8 | 19.6 | - |
| +SFT | 54.2 | 25.0 | 35.7 | 38.5 | 70.9 | 86.1 | 39.7 | 49.3 |
| +DPO | 52.8 | 9.0 | 16.6 | 35.0 | 83.5 | 87.5 | 37.9 | 49.1 |
| JetMoE-2B-9B | 45.6 | 43.0 | 37.2 | 54.6 | - | 68.2 | 20.0 | - |
| +SFT | 46.1 | 53.5 | 35.6 | 64.8 | 69.3 | 55.6 | 30.5 | 50.4 |
| DeepSeek-3B-16B | 37.7 | 18.5 | 39.4 | 48.3 | - | 65.9 | 13.5 | - |
| +Chat | 48.5 | 46.5 | 40.8 | 70.1 | 74.8 | 85.6 | 32.3 | 57.0 |
| Qwen1.5-3B-14B | 60.4 | 13.5 | 27.2 | 60.2 | - | 73.4 | 20.9 | - |
| +Chat | 58.9 | 55.5 | 21.3 | 59.7 | 83.9 | 85.6 | 36.2 | 57.3 |
| OLMoE-1B-7B | 49.8 | 3.0 | 33.6 | 22.4 | - | 59.7 | 16.6 | - |
| +SFT | 51.4 | 40.5 | 38.0 | 51.6 | 69.2 | 84.1 | 43.3 | 54.0 |
| +DPO | 51.9 | 45.5 | 37.0 | 54.8 | 84.0 | 82.6 | 48.1 | 57.7 |

##### After pretraining

In Table 4 we benchmark OLMoE-1B-7B on common downstream tasks. We find that OLMoE-1B-7B performs best among models that use less than 2B active parameters, making it the most economical option for many use cases of LMs. For larger budgets, Qwen1.5-3B-14B has stronger performance but has more than double the active and total parameters than OLMoE-1B-7B. We find that despite requiring 6–7 less compute per forward pass, OLMoE-1B-7B outperforms some dense LMs with 7B parameters such as Llama2-7B [], but falls short of others like Llama3.1-8B []. Figure 1 compares MMLU performance with active parameters, a proxy for the value of a model given its cost, of OLMoE-1B-7B and other LMs. OLMoE-1B-7B is the state of the art in its cost regime.

##### After adaptation

In Table 5, we benchmark our instruction (SFT) and preference (DPO) tuning of OLMoE-1B-7B. SFT improves our model on all tasks measured. We observe a $>$10 gain on GSM8k, likely due to our inclusion of additional math data to account for the relatively small amounts of math data during pretraining (§2). DPO helps on most tasks, especially AlpacaEval which aligns with findings from prior work [, , ]. Our DPO model, which we refer to as OLMoE-1B-7B-Instruct, has the highest average among all models benchmarked. We find it to outperform the chat version of Qwen1.5-3B-14B despite Qwen having $>$2 more parameters and its pretrained model outperforming OLMoE-1B-7B in Table 4. The 84% score on AlpacaEval also outperforms much larger dense models on the leaderboard,(^2^22[https://tatsu-lab.github.io/alpaca_eval/](https://tatsu-lab.github.io/alpaca_eval/)) such as Llama2-13B-Chat [].

## 4 Experimenting with Alternative Design Choices

In this section, we present pretraining and adaptation experiments that have led to OLMoE-1B-7B. We group them into experiments on settings specific to Mixture-of-Experts (§4.1), experiments on settings applicable to both dense LMs and MoEs (§4.2), and adaptation experiments (§4.3). In pretraining experiments, we often use MMLU Var, a version of MMLU [] with varying few-shots and a different format that provides signal earlier during training. We describe our full evaluation setup in Appendix C and provide additional experiments in Appendix F. Each experiment links to a Weights & Biases report with more validation and downstream results, and the full configurations of the runs. To isolate the impact of changes and minimize confounders, we vary only one hyperparameter for each experiment. Nevertheless, due to the large number of hyperparameters, some results may change under different configurations and we cannot guarantee the correctness of each of our hyperparameter choices. Models are not comparable across different experiments, as we vary the base model to incorporate successful findings.

### 4.1 MoE-specific Pretraining Settings

#### 4.1.1 Mixture-of-Experts vs. Dense

Figure: Figure 4: MoE vs. Dense. We train a 1.3B parameter dense model and a 1.3B active, 6.9B total parameter MoE model, each on 128 H100 GPUs. Apart from MoE-related changes, we train both with the same configuration for 130B tokens. The MoE contains 64 experts out of which 8 are activated with an FFN dimension of 1,024, while the dense model has an FFN dimension of 8,192. Thus both have the same number of active parameters. Top: The MoE reaches the final dense performance with 3 fewer tokens (or FLOPs, as both have the same active parameters ignoring the trivial router parameters). Bottom: Due to some memory overhead, this equates to 2 faster training. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-MoE-vs-Dense--Vmlldzo4OTM0Mjkx](https://wandb.ai/ai2-llm/olmoe/reports/Plot-MoE-vs-Dense--Vmlldzo4OTM0Mjkx)
Refer to caption: x6.png

Prior work reports various speed-ups of MoEs over dense models: Artetxe et al. [] report that MoEs require 2–4 less compute to match dense models, MoMa [] exhibits 2.6 FLOP savings for language tasks, Arctic [] yields 4 FLOP savings but for very different dense and MoE configurations, and Switch Transformers [] train 2-7 faster with MoEs but for encoder-decoder models while the other works study decoder-only LMs [].

In Figure 4, we compare MoEs and dense models in a controlled setup. We find that our MoE reaches the performance of the dense model with 3 fewer tokens equivalent to 3 less compute measured in FLOPs. However, due to the additional memory overhead of training the MoE with its 7B total parameters, it processes fewer tokens per second than the dense model (23,600 tokens per second per GPU for the MoE vs. 37,500 for dense). Thus, in terms of training time, it reaches the performance of the dense model only 2 faster. There are likely optimizations possible that would bring the speed-up closer to the 3 token speed-up, which we leave to future work. Based on these results, we select an MoE configuration with 6.9B total and 1.3B active parameters matching OLMo-7B in total and OLMo-1B in active parameter count, respectively.

#### 4.1.2 Expert Granularity

Figure: Figure 5: Expert granularity. We vary the number of experts in tandem with the FFN dimension to ensure that active and total parameters and thus compute cost remain the same. For example, for 64 experts, the FFN dimension is 1,024 and 8 experts are activated, while for 32 experts it is 2,048 with 4 activated experts. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Granularity--Vmlldzo4OTIxOTE4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Granularity--Vmlldzo4OTIxOTE4)
Refer to caption: x7.png

Dai et al. [] propose to use small fine-grained experts to allow more combinations of experts and thus make the model more flexible. For example, the Mixtral model [] uses the common configuration of 8 experts per layer, 2 of which are activated. This allows for $\binom{8}{2}=28$ combinations per layer. By halving the size of each expert and therefore doubling the number of experts to maintain the same compute and parameter budget, we can increase the possible combinations to $\binom{16}{4}=1,820$. Krajewski et al. [] investigate compute-optimal granularity configurations finding that higher compute budgets warrant more granular experts.

In Figure 5, we observe that more granular experts improve training loss, validation loss, and downstream performance. The 8-expert configuration uses 1 active expert, which yields $\binom{8}{1}=8$ combinations. By quartering the size of each expert but increasing the number to 32 with 4 active ones ($\binom{32}{4}=35,960$ combinations), we observe an improvement of around 10% on HellaSwag and MMLU at around 130 billion tokens. However, we find that there are diminishing returns to granularity. The additional increase to 64 experts with 8 active ones ($\binom{64}{8}=4,426,165,368$ combinations) improves downstream metrics by a smaller amount of 1–2%. For our OLMoE-1B-7B compute budget(^3^33Approximated via $6*N*D$ [], where $N$ are active parameters (1B) and $D$ are training tokens (5T).) of $3\times 10^{22}$, Krajewski et al. [] predict an optimal number of experts of 256 ($G=32$ in their paper). However, their predictions are for compute-optimal models [, ], while we train for 5T tokens, which is orders of magnitude beyond what would be conventionally considered optimal for our model size. Thus, their predictions may not extend to our setup, and we stick with 64 experts for OLMoE-1B-7B, also due to the diminishing returns in Figure 5.

#### 4.1.3 Shared Experts

Figure: Figure 6: Shared experts. Both setups have the same number of active and total parameters and use the same number of FLOPs. 4 of the 32 routed experts are activated, while it is 3 for the 31 routed experts of the other model, as it has 1 always-active shared expert. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Expert-sharing--Vmlldzo4OTIyMjQz](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Expert-sharing--Vmlldzo4OTIyMjQz)
Refer to caption: x8.png

Dai et al. [] propose training with a shared/fixed expert that is always used in addition to the routed experts. The intuition is to encourage the shared expert to learn common information and allow the other routed experts to learn more specialized knowledge. This should reduce redundancy among experts and thus lead to a better model as it can store more total information.

In Figure 6, we benchmark having a single shared and a single routed expert versus two routed experts. While both settings lead to similar performance, sharing an expert performs slightly worse. Sharing an expert removes flexibility from the model and thus goes against the findings in §4.1.2 suggesting that allowing for more expert combinations improves performance. Specifically, the two models in Figure 6 have $\binom{32}{4}=35,960$ and $\binom{31}{3}=4,495$ possible combinations per layer. Thus, removing one of the routed experts and turning it into a shared one eliminates almost 90% of possible combinations. This likely acts as a counterforce to the potential benefits of isolating common knowledge in a shared expert. Based on these results, we do not use shared experts in OLMoE-1B-7B, but we do think that there is merit to the idea of experts that are activated more often or even always. However, rather than enforcing this behavior via a shared expert, we believe that it should be learned by the model. This is difficult with current setups due to the necessity of a load balancing loss (§4.1.6) penalizing the model if tokens are not distributed equally among experts. Potential future work can explore removing the load balancing loss to allow for more flexible usage of experts.

#### 4.1.4 Expert Choice vs. Token Choice

Figure: Figure 7: Expert choice (EC) vs. token choice (TC). Both models have an 8-expert MoE in every 2nd layer. For TC, 2 experts are activated per token, while for EC the capacity factor is 2. Thus, both models use the same number of active parameters. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-EC-vs-TC--Vmlldzo4MzkzMDM3](https://wandb.ai/ai2-llm/olmoe/reports/Plot-EC-vs-TC--Vmlldzo4MzkzMDM3)
Refer to caption: x9.png

The MoE router determines which experts process each input token (§2). There are two common types []: expert choice (EC) [] and token choice (TC) []. For EC, each expert selects a fixed number of tokens from the incoming sequence. By design, this leads to each expert processing the same number of tokens. This is the main benefit of EC as it ensures perfect load balance, which improves training throughput and removes the need for a load balancing loss. The main downside of EC is that it is not easily usable for autoregressive generation where a single token is processed at each step rather than the entire sequence in one []. Another potential downside is that EC can lead to token dropping, where some tokens are not selected by any expert, which can hurt performance []. At the same time, it can lead to some tokens being processed by multiple experts, which could also be beneficial as it allows the model to allocate more compute to some tokens []. For TC, each token selects a fixed number of experts. This can lead to many tokens choosing the same expert, hurting training efficiency. Therefore it is common to use TC with a load balancing loss [] to encourage equal distribution.

In Figure 7, we benchmark EC and TC. We find that TC outperforms EC for the same token budget for all tasks depicted as well as other tasks like PIQA, SciQ, etc. which we report at [https://wandb.ai/ai2-llm/olmoe/reports/Plot-EC-vs-TC--Vmlldzo4MzkzMDM3](https://wandb.ai/ai2-llm/olmoe/reports/Plot-EC-vs-TC--Vmlldzo4MzkzMDM3). While Zhou et al. [] find EC to be better, our configuration slightly differs in that we use dropless MoEs [] with a load balancing loss. Thus, our TC variant is expected to perform better than the TC variant in Zhou et al. []. We confirm findings that EC runs around 20% faster at 29,400 tokens per second per device versus 24,400 for TC []. EC may be more beneficial in a multimodal setup [] as dropping noisy image tokens is likely less harmful than text tokens. Thus, while we stick with TC for this release of OLMoE, we may revisit EC for future multimodal models.

#### 4.1.5 Sparse Upcycling

Figure: Figure 8: Sparse upcycling. We upcycle OLMo-1B (0724) at 2T tokens into an MoE with 8 total experts of which 2 are activated and train it for an additional 610 billion tokens. We compare it to a model trained from scratch for 610 billion tokens. Except for this difference, both models use the same config, which includes some suboptimal settings that contribute to the instability, such as no QK-Norm (§4.2.5) and no truncated normal init (§4.2.2). More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Scratch-vs-Upcycle--Vmlldzo4NDIyOTc4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Scratch-vs-Upcycle--Vmlldzo4NDIyOTc4)
Refer to caption: x10.png

Komatsuzaki et al. [] propose turning a dense model into a Mixture-of-Experts model via sparse upcycling: (1) The dense MLP is cloned for each desired expert to constitute MoE layers. (2) A newly initialized router is added in front of each MoE layer. (3) Pretraining continues with the new model so that the cloned MLPs can gradually specialize in different things and the router can be learned. They find that the upcycling approach maintains a performance advantage over a language model trained from scratch for up to 120% of the compute budget of the original dense checkpoint that the sparse model was upcycled from. For example, if sparsely upcycling a 1.3B parameter model at 2 trillion tokens then only at 2.4 trillion tokens should an MoE trained from scratch catch up with the upcycled model. That is, the sparsely upcycled model would have been trained for another 400 billion tokens, thereby saving the equivalent of up to 2T tokens of compute. Other works such as MiniCPM [], Qwen2 [] and reportedly Mixtral [, ] have adopted sparse upcycling but only share limited information about their configuration.

In Figure 8, we compare sparse upcycling OLMo-1B (0724) [] with training an MoE from scratch. We find that after 500B tokens, an otherwise equivalent MoE trained from scratch already catches up with the upcycled model, both on the metrics in Figure 8 and our additional metrics at [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Scratch-vs-Upcycle--Vmlldzo4NDIyOTc4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Scratch-vs-Upcycle--Vmlldzo4NDIyOTc4). At around 600B tokens, the MoE from scratch starts outperforming the upcycled MoE. Thus, it only requires 25% of the compute budget of the original dense model to catch up as opposed to the 120% reported in Komatsuzaki et al. []. However, they use expert choice routing and study encoder-decoder models []. Meanwhile, we use token choice routing (§4.1.4) and decoder-only models (§2). Further, we upcycle a model that has already been significantly overtrained [], i.e., a 1B model trained for 2T tokens. Its parameters are likely already in a very optimal range for a dense model, which may limit the amount of additional exploration possible after upcycling. This motivates us to experiment with adding noise to the upcycled weights outlined in Appendix F, but we do not find it to lead to better performance. A large disadvantage of upcycling is that the upcycled MoE is constrained by some hyperparameters of the dense model. Specifically, OLMo-1B (0724) was trained without QK-Norm and normal initialization, both of which hurt stability in our experiments (§4.2.5, §4.2.2). While it may be possible to simply add new QK-Norms and train them from scratch similar to the new router layer trained from scratch, it is impossible to change the initialization of the original dense model when upcycling it. Thus, as we want to change these hyperparameters and also train OLMoE-1B-7B for around 250% of the compute budget of the dense model (5T vs. 2T tokens), we do not use upcycling.

#### 4.1.6 Load Balancing Loss

Figure: Figure 9: Impact of applying a load balancing loss (LBL). The training loss plot excludes the load balancing loss for both models. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-LBL-vs-No-LBL--Vmlldzo4OTkyNDg4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-LBL-vs-No-LBL--Vmlldzo4OTkyNDg4)
Refer to caption: x11.png

Figure: Figure 10: Expert assignment during training when using or not using a load balancing loss for the first MoE layer. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-LBL-vs-No-LBL--Vmlldzo4OTkyNDg4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-LBL-vs-No-LBL--Vmlldzo4OTkyNDg4)
Refer to caption: x12.png

Shazeer et al. [] propose the load balancing loss to penalize the model if it is unbalanced, i.e., if it routes all tokens to only a few experts. This is based on the observation that without such penalty, models tend to update only a select few experts in each layer [, ]. To compute the load balancing loss ($\mathcal{L}_{\textit{LB}}$) we multiply the fraction of tokens $f_{i}$ routed to one expert $E_{i}$ with the total routing probability $P_{i}$ allocated to $E_{i}$ for one batch and sum it across the number of experts $N_{E}$:

$$ $\mathcal{L}_{\textit{LB}}=N_{E}\cdot\tsum\slimits@_{i=1}^{N_{E}}f_{i}\cdot P_{i}$ (3) $$

The loss is further scaled by $N_{E}$ and a loss weight $\alpha$ (see Equation 2), which is an optional weight to determine the magnitude of the loss commonly set to 0.01 [, ]. We do not experiment with changing the weight of 0.01.

In Figure 9 we investigate the performance impact of using the auxiliary load balancing loss. We find that across training loss and validation losses, using the load balancing loss leads to better performance even after only a few billion tokens. We still measure the load balancing loss even when it is not used (“No LBL”) and find that while it spikes initially, it slowly decreases over the next few billion tokens. This behavior is also visible in Figure 10 (left), where initially all tokens in the first layer are assigned to the 6th expert (pink). Eventually, the model also starts assigning some tokens to the 1st expert (yellow). However, all other experts remain largely flat and are thus “dead weights” that take up GPU memory but are not used. Given these results, we use the auxiliary load balancing loss with a weight of 0.01 following prior work [, ]. However, getting rid of the load balancing loss is an important direction for future research as it constrains the flexibility of the model by forcing it to use all experts approximately equally. This could prevent the experts from specializing in certain data domains and may be a reason prior work has failed to find strong evidence of expert specialization [, ].

#### 4.1.7 Router Z-loss

Figure: Figure 11: Router z-loss. We compare adding router z-loss with a loss weight of 0.001 versus no additional z-loss. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Zloss-vs-none--Vmlldzo4NDM4NjUz](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Zloss-vs-none--Vmlldzo4NDM4NjUz)
Refer to caption: x13.png

Zoph et al. [] propose the router z-loss to improve both the stability and quality of MoE models. This auxiliary loss penalizes large logits coming into the gating network. Such large logits can lead to numeric overflows in the large matrix multiplications happening in the MoE layer. It is computed by exponentiating the logits $x_{j}$ right before the router layer summed across the number of experts $N_{E}$ and averaged across the batch $B$, thereby making larger logits lead to a larger loss:

$$ $\mathcal{L}_{\textit{RZ}}(x)=\frac{1}{B}\cdot\tsum\slimits@_{i=1}^{B}\left(\log\tsum\slimits@_{j=1}^{N_{E}}\exp({x_{j}^{(i)}})\right)^{2}$ (4) $$

The loss is further multiplied with an optional loss weight, $\beta$ (see Equation 2), to determine the magnitude of the loss commonly set to 0.001 [, ]. We do not experiment with changing the weight of 0.001.

In Figure 11, we confirm that across training loss, validation loss, and downstream performance adding the router z-loss improves stability (less spikes) and quality (lower loss and higher downstream performance). Thus, despite it reducing throughput by 2% we use the router z-loss for OLMoE-1B-7B with a weight of 0.001 as in Zoph et al. [].

### 4.2 General Pretraining Settings

#### 4.2.1 Dataset Experiments

Figure: Figure 12: OLMoE-Mix vs. Dolma 1.7. We compare our data mix described in §2 with Dolma 1.7 used to train prior OLMo models. Lower training loss does not mean that one dataset is better, but rather suggests which dataset is easier for the model to learn. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Dolma-1-7-vs-Dolma-OLMoE--Vmlldzo4OTIxNTg5](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Dolma-1-7-vs-Dolma-OLMoE--Vmlldzo4OTIxNTg5)
Refer to caption: x14.png

Li et al. [] release the DCLM-Baseline dataset and establish that it leads to better language models than Dolma 1.7 and other datasets as measured on common benchmarks like MMLU []. This motivates us to mix their DCLM dataset with some components from Dolma 1.7 that we deem to be high-quality; see §2. In Figure 12, we compare our mix, OLMoE-Mix, with Dolma 1.7 in a controlled setup. We find that OLMoE-Mix leads to clear gains on all three downstream metrics, especially MMLU. DCLM-Baseline has been created through a series of dataset ablations targeting MMLU and other downstream metrics, which explains these results. We also compare adding Reddit and FLAN to our mix as detailed in Appendix F, but do not find consistent performance gains. We do not have a strong intuition for why adding these datasets does not help and a more automatic approach to dataset mixing may be desirable for future iterations [, ]. We pretrain using our mix of DCLM-Baseline and Dolma 1.7 dubbed OLMoE-Mix.

#### 4.2.2 Initialization

Figure: Figure 13: Initialization. We compare a normal initialization with a standard deviation (std) of 0.02 with a truncated normal initialization with a maximum (minimum) cut-off of 0.06 (–0.06) corresponding to three stds (30.02). More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Init--Vmlldzo4NDIzMzM5](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Init--Vmlldzo4NDIzMzM5)
Refer to caption: x15.png

Few prior works on Mixture-of-Experts share their initialization strategy. Even the most open MoEs prior to this work, JetMoE [] and OpenMoE [], do not mention their initialization scheme. For DeepSeekMoE [] and DeepSeekV2 [], the authors share that they use a normal initialization with a standard deviation (std) of 0.006. For dense language models, a normal initialization with an std of 0.02 has been commonly used as popularized by Shoeybi et al. [].

In Figure 13, we find a truncated normal initialization leads to more stable training and better performance than a regular normal initialization. The difference between the two initializations only becomes clear at around 450 billion tokens, where the model with the normal initialization starts to diverge. This is despite both models using the same configuration except for the difference in weight initialization. Having to train for hundreds of billions of tokens until an experiment provides a clear signal is one of the key challenges of pretraining ablations. We use the truncated normal initialization for OLMoE-1B-7B.

#### 4.2.3 RMSNorm

Figure: Figure 14: Non-parametric layer normalization vs. RMSNorm. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-LN--Vmlldzo4NDQyMTAz](https://wandb.ai/ai2-llm/olmoe/reports/Plot-LN--Vmlldzo4NDQyMTAz)
Refer to caption: x16.png

Figure: Figure 15: Decaying the RMSNorm parameters. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Decay-LN--Vmlldzo4NDQ1NDYy](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Decay-LN--Vmlldzo4NDQ1NDYy)
Refer to caption: x17.png

Figure: Figure 16: Total norm of the gradients when training with RMS or non-parametric normalization. We increase the logging interval of the RMS run at 75B tokens, hence its change in thickness.
Refer to caption: x18.png

OLMo [] uses non-parametric layer normalization [], mainly as it is significantly faster than the commonly used RMSNorm [, ]. This is an unusual choice as most LMs use RMSNorm, such as the Llama [, , ], Gemma [, ], and Qwen [, ] model families.

In Figure 14, we observe that replacing the non-parametric layer normalization in OLMo with a parametric RMSNorm leads to better performance. This is likely because the non-parametric layer normalization leads to a large number of spikes in the gradients as seen in Figure 16. We clip gradients at 1.0, which prevents these spikes from leading to very large and potentially disruptive parameter updates. However, the clipped gradients may still harm the performance of the model as they are no longer the true gradients. Thus, despite RMSNorm lowering our training throughput by 15%, we train our final model with RMSNorm. We include the RMSNorm parameters in weight decay as we find that it performs slightly better (Figure 15) even though it is common practice to exclude them.(^4^44[https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025](https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025))

#### 4.2.4 Decaying Embedding Parameters

Figure: Figure 17: Decaying the embedding parameters. More results, logs, and configurations: [https://api.wandb.ai/links/ai2-llm/3h22onp5](https://api.wandb.ai/links/ai2-llm/3h22onp5)
Refer to caption: x19.png

Similar to the RMSNorm parameters (§4.2.3), embedding parameters are commonly excluded from weight decay.(^5^55[https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025](https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025)) In Figure 17 we find that whether or not they are decayed has only a minor impact on performance, with decaying being slightly better. Thus for simplicity, we weight decay all parameters in OLMoE-1B-7B including embedding and RMSNorm.

#### 4.2.5 QK-Norm

Figure: Figure 18: Query-Key layer normalization (QK-Norm). Both models use non-parametric layer normalization. QK-Norm corresponds to additional layer normalization of the query and key projections. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-QKNorm-vs-none--Vmlldzo4NDIzMzE2](https://wandb.ai/ai2-llm/olmoe/reports/Plot-QKNorm-vs-none--Vmlldzo4NDIzMzE2)
Refer to caption: x20.png

Some works have reported stability improvements from adding layer normalization after the query and key projections (“QK-Norm”) [, , ]. QK-Norm can prevent the subsequent attention operation from leading to very large logits that may lead to numeric overflows and destabilize the network, especially when training in low precision. Like layer normalization at other places in the model, the QK-Norm could be non-parametric or use the parametric RMSNorm (§4.2.3).

In Figure 18, we compare using QK-Norm with no normalization after the query and key projections. We find that QK-Norm leads to some stability and performance improvements. We perform this experiment with non-parametric layer normalization as used in OLMo [], while we used parametric RMS layer normalization [] for OLMoE-1B-7B (§4.2.3). To ensure the benefit of QK-Norm is not an artifact of comparing with non-parametric layer normalization, we run another experiment with RMS layer normalization and still find QK-Norm to lead to slightly better training loss and to prevent a large grad norm spike.(^6^66[https://wandb.ai/ai2-llm/olmoe/reports/Plot-QKNorm-revisited--Vmlldzo4NTc2NTIz](https://wandb.ai/ai2-llm/olmoe/reports/Plot-QKNorm-revisited--Vmlldzo4NTc2NTIz)) Thus, we use QK-Norm for OLMoE-1B-7B despite it reducing throughput by almost 10%.

#### 4.2.6 AdamW Epsilon

Figure: Figure 19: AdamW epsilon. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-AdamW-eps--Vmlldzo4NDc5MDg0](https://wandb.ai/ai2-llm/olmoe/reports/Plot-AdamW-eps--Vmlldzo4NDc5MDg0)
Refer to caption: x21.png

Groeneveld et al. [] use an epsilon (“eps”) value of 1E-05 in the AdamW optimizer for training OLMo. A larger eps value leads to smaller steps of the optimizer but can be more stable [].

In Figure 19, we find that decreasing eps to the recommended default of 1E-08 [] significantly improves performance while the run remains stable. Thus, we set eps to 1E-08 for our final run.

### 4.3 Adaptation Settings

**Table 6: Load balancing loss (Equation 3) over a subset of the respective corpora prior to scaling with the load balancing loss weight $\alpha$. While we use load balancing loss during pretraining, we do not use it during SFT.**
| Data ($\downarrow$) | OLMoE-1B-7B |  |
| --- | --- | --- |
| After pretraining | After SFT |  |
| SFT data | 12.22 | 12.16 |
| Github | 13.85 | 14.85 |
| Wikipedia | 14.48 | 14.24 |
| C4 | 9.09 | 9.13 |

We experiment with small design choices for adaptation using our evaluation setup described in Appendix C. (1) Auxiliary losses: Zoph et al. [] find that using the auxiliary load balancing loss (§4.1.6) during regular finetuning leads to small performance gains. For instruction tuning, however, Shen et al. [] do not find conclusive evidence in favor of using the load balancing or router z-loss with only small differences in performance, both in support of and against the auxiliary losses. In Table 7 we display experiments with the load balancing loss during adaptation and find that not using it leads to better performance (54.0 vs. 52.8 after instruction tuning (SFT) and 57.7 vs. 57.1 after preference tuning (DPO)). One potential problem of deactivating the load balancing loss is that it may harm balance among experts and turn some into dead weights as observed during pretraining in §4.1.6. However, when measuring the load balancing loss in Table 6 on our SFT data (§2), we find that the loss actually decreases slightly during SFT (12.16 vs. 12.22).
This is likely because which experts certain tokens get routed to is determined early during pretraining, as we find later in the analysis section (§5.1). We also visualize the activation patterns of experts of the model after pretraining, and the models after SFT and DPO trained without load balancing in Appendix G (Figure 33) finding that the distribution remains around the same. Thus, as our models adapted without load balancing perform better and we find it not to impact routing substantially, we do not use load balancing during adaptation. (2) Annealing checkpoint: We also experiment with using the checkpoint pre-annealing (§2) for adaptation and find the checkpoint post-annealing leads to better performance (53.8 vs. 54.0 after SFT and 56.3 vs 57.7 after DPO), thus we use the post-annealing checkpoint. (3) Preference algorithm: Since the release of DPO (Direct Preference Optimization) [], a variety of preference algorithms have been proposed [, , ]. We experiment with KTO [] and find that it matches DPO in Table 7 for our setup (Appendix B). While we release both models, we use DPO for our final OLMoE-1B-7B-Instruct model, as it scores higher on AlpacaEval, which has a smaller chance of data contamination than our other benchmarks [].

**Table 7: Adaptation experiments of OLMoE-1B-7B. We compare using the pretrained checkpoint prior to annealing for adaptation, using the checkpoint after the additional 100B tokens of annealing, and using the checkpoint after the additional 100B tokens of annealing and with load balancing loss (§4.1.6) during adaptation. We apply DPO/KTO to the respective SFT model.**
|  |  |  |  | Human- | Alpaca- |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Task () | MMLU | GSM8k | BBH | Eval | Eval 1.0 | XSTest | IFEval | Avg |
| Setup () | 0-shot | 8-shot CoT | 0-shot | 0-shot | 0-shot | 0-shot | 0-shot | 0-shot |
| Metric () | EM | EM | EM | Pass@10 | %win | F1 | Loose Acc |  |
| OLMoE-1B-7B<br>w/o annealing | 49.0 | 2.0 | 31.5 | 18.9 | - | 62.1 | 18.5 | - |
| +SFT | 50.2 | 43.0 | 35.6 | 55.5 | 68.9 | 83.8 | 39.7 | 53.8 |
| +DPO | 50.9 | 36.0 | 35.8 | 58.8 | 81.7 | 83.2 | 47.9 | 56.3 |
| OLMoE-1B-7B | 49.8 | 3.0 | 33.6 | 22.4 | - | 59.7 | 16.6 | - |
| +SFT | 51.4 | 40.5 | 38.0 | 51.6 | 69.2 | 84.1 | 43.3 | 54.0 |
| +DPO | 51.9 | 45.5 | 37.0 | 54.8 | 84.0 | 82.6 | 48.1 | 57.7 |
| +KTO | 51.2 | 45.5 | 34.1 | 57.1 | 81.6 | 86.6 | 47.5 | 57.7 |
| +SFT<br>(load balancing) | 50.9 | 36.5 | 35.7 | 52.4 | 66.9 | 84.8 | 42.3 | 52.8 |
| +DPO<br>(load balancing) | 51.1 | 42.5 | 39.3 | 55.6 | 82.9 | 82.1 | 46.0 | 57.1 |

## 5 MoE Analysis

By advancing open and cost-efficient models (§1), OLMoE-1B-7B enables new research into LMs and MoEs. Making use of our released intermediate checkpoints, data, and code, we define and analyze four properties specific to MoEs: Router saturation (§5.1), Expert co-activation (§5.2), Domain specialization (§5.3), and Vocabulary specialization (§5.4).

### 5.1 Router Saturation

Figure: Figure 20: Router saturation during pretraining measured on a random 0.5% of the C4 validation data. We compute saturation by comparing the routing to the top-$k$ experts at four intermediate checkpoints (1, 10, 20, and 40% of pretraining) to the final pretraining checkpoint (Equation 5).
Refer to caption: x22.png

We define router saturation as the proportion of expert activations at some intermediary checkpoint at time $t$ that matches the expert IDs activated at some final checkpoint over the same dataset:

$$ $\text{Router Saturation}(t)=\frac{1}{N}\tsum\slimits@_{i=1}^{N}\frac{|\mathcal{E}_{i}^{(t)}\cap\mathcal{E}_{i}^{(T)}|}{k},$ (5) $$

where:

- •
$N$: The total number of tokens in the dataset.
- •
$k$: The number of top-$k$ experts activated per input token. While we train with $k=8$ (§2), we also analyze $k=1$ by only looking at the expert with the highest routing probability.
- •
$\mathcal{E}_{i}^{(t)}$: The set of $k$ experts activated for the $i$th token at the $t$th checkpoint.
- •
$\mathcal{E}_{i}^{(T)}$: The set of $k$ experts activated for the $i$th token at the final checkpoint $T$.
- •
$|\mathcal{E}_{i}^{(t)}\cap\mathcal{E}_{i}^{(T)}|$: The number of common experts activated for the $i$th token between the $t$th and final checkpoints.

Router saturation thus corresponds to whether the router weights are still learning which expert will process certain data. A value of 100% indicates that the router at the intermediate checkpoint will route to the same experts as the final checkpoint router. However, even at 100% saturation the router weight can still change and adapt the exact router probability for each expert. These probabilities are used to scale the output of the respective expert in the model. For OLMoE-1B-7B with its 64 experts, random routing equals a saturation of $1/64=1.6\%$ for $k=1$ and $8/64=12.5\%$ for $k=8$.

In Figure 20 we find that after 1% of pretraining (5000 steps or 20B tokens), up to 60% of routing to the top-8 activated experts has already saturated (right). Thus the model already uses the same 8 experts for given input data as it will at the end of pretraining. This early saturation aligns with prior work []. At 40% of pretraining, saturation reaches up to 80%. However, which top-1 expert has the highest routing probability saturates slower (left). We find that routing in later layers saturates earlier during pretraining. Layer 0 is an outlier saturating significantly more slowly than other layers. Dai et al. [] do not use an MoE in the first layer as they find that load balancing converges more slowly for the first layer. This is likely linked to our findings on saturation. Because routing in the first layer saturates slower, the experts that certain input data get routed to frequently change. These changes may lead to one expert suddenly getting significantly more data than others thereby impairing load balancing. We are excited about future work further investigating what happens in the first layer by building on our open release.

### 5.2 Expert Co-activation

Figure: Figure 21: Co-activation among experts of OLMoE-1B-7B on a random 0.5% of the C4 validation data. We display the 32 experts with the highest maximum co-activation score via their expert IDs on the x- and y-axis.
Refer to caption: x23.png

We define expert co-activation as the proportion of times two specific experts, $E_{i}$ and $E_{j}$, are simultaneously activated out of the total number of activations of one of those experts:

$$ $\text{Expert co-activation}(E_{i},E_{j})=\frac{N_{E_{i},E_{j}}}{N_{E_{i}}},$ (6) $$

where:

- •
$E_{i}$: The first expert.
- •
$E_{j}$: The second expert.
- •
$N_{E_{i},E_{j}}$: The number of times experts $E_{i}$ and $E_{j}$ are activated together.
- •
$N_{E_{i}}$: The total number of times expert $E_{i}$ is activated.

A co-activation of 100% indicates that if $E_{i}$ is activated, $E_{j}$ is also always activated. A value of 0% indicates that the experts never co-occur. If multiple expert pairs have high co-activation, it may suggest that these experts could be merged, benefiting less from keeping them separate. In a distributed setup, we could place highly co-activated experts on the same device to reduce communication costs during model inference.

In Figure 21, we find that there is no strong co-activation among experts in one layer, with only few exceptions. This may indicate that there is little redundancy across different experts. Overall, layers 7 and 15 show similar co-activation patterns with several groups of 3 or 2 experts that tend to get activated together. We investigate tokens that activate these experts in §5.4. Further, in Appendix G (Figure 35), we investigate whether experts across layers, rather than within one layer, tend to process tokens together.

### 5.3 Domain Specialization

Figure: Figure 22: Domain specialization of OLMoE-1B-7B (top) vs. Mixtral-8x7B (bottom). We visualize how often tokens from different domains get routed to the 64 (OLMoE) or 8 (Mixtral) experts at the end of pretraining. We consider tokens routed to any of the $k=8$ (OLMoE) or $k=2$ (Mixtral) active experts (Equation 7). Horizontal gray lines correspond to random chance or uniform routing (8/64=12.5% per expert for OLMoE-1B-7B with 8 active out of 64 total experts per layer and 2/8=25% for Mixtral with 2 active out of 8 total experts per layer). See Figure 34 for $k=1$ results.
Refer to caption: x26.png

We define domain specialization as the proportion of tokens from a particular domain $D$ that get routed to a particular expert $E_{i}$:

$$ $\text{Domain specialization}(E_{i},D)=\frac{N_{E_{i},D}^{(k)}}{N_{D}},$ (7) $$

where:

- •
$E_{i}$: The $i$th expert in the model.
- •
$D$: The domain from which the data originates.
- •
$k$: The number of experts considered (e.g., $k=8$ means considering the top 8 experts with the highest routing probabilities).
- •
$N_{E_{i},D}^{(k)}$: The number of tokens from domain $D$ for which $E_{i}$ is among the top-$k$ selected experts.
- •
$N_{D}$: The total number of tokens from domain $D$ processed by the MoE.

Domain specialization thus refers to the specialization of expert $E_{i}$ to domain $D$. A value of 100% indicates that all data from that domain is routed to $E_{i}$, whereas 0% indicates the expert is never used for that domain and can be removed from the model without affecting performance in that domain.

In Figure 22 (top) we find many examples of experts that are activated significantly above or below random chance for specific domains. E.g., for arXiv, which has a very specific distribution with lots of scientific text, the first expert in layer 0 is nearly 100% specialized. This suggests that there is little redundancy in the knowledge of the experts in OLMoE-1B-7B, as they specialize in different kinds of data. GitHub and arXiv are often activated together in layer 7, which we explore further in §5.4. For generic domains, such as C4 [], which is a web crawl containing various kinds of data, expert activations in OLMoE-1B-7B are much more balanced. This highlights that the load balancing (§4.1.6) works as intended and the model makes proper use of all experts for generic data. Mixtral-8x7B [] in Figure 22 (bottom), however, exhibits little domain specialization across both unique and generic domains. Experts are activated close to the uniform routing baseline for all layers and domains. Thus, there may be more redundancy across experts in Mixtral, as they likely contain similar knowledge. We hypothesize that this is due to Mixtral being upcycled from Mistral []. The initialization from a dense model may limit the amount of possible specialization in the experts as they all start from the same local optimum. This is likely why training from scratch eventually outperforms upcycling in our pretraining experiments (§4.1.5).

### 5.4 Vocabulary Specialization

Figure: Figure 23: Vocabulary specialization of OLMoE-1B-7B across layers and experts. To compute vocabulary specialization per layer (left) we average the specialization of each expert in that layer. Dashed lines (right) correspond to the average of layer 7 as depicted left. We display the first 32 experts out of 64. This plot is for $k=1$ (Equation 8) and we provide $k=8$ and a comparison with Mixtral-8x7B in Appendix G.
Refer to caption: x28.png

**Table 8: Vocabulary specialization in the 7th layer of OLMoE-1B-7B. We use $k=1$ (Equation 8) and a random 0.5% of the C4 validation data excluding token IDs with $<$10 appearances.**
| Expert ID | Input token IDs | Predicted output token IDs |
| --- | --- | --- |
| 27 | (100%) (100%) ^3 (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) | (100%) § (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) |
| 58 | (“ (100%) (” (100%) ‘ (94%) ’ (92%) “ (92%) ( (92%) ” (90%) ’ (89%) “ (88%) $ (87%) [ (87%) £ (86%) | such (100%) 486 (100%) see (95%) which (91%) driving (91%) UK (90%) who (88%) including (88%) normal (88%) |
| 7 | Him (100%) inde (100%) Jesus (98%) God (90%) pray (81%) Holy (80%) Quran (80%) God (77%) Lord (76%) glory (75%) Spirit (66%) Christ (65%) | rella (100%) Him (94%) sin (90%) prince (80%) glory (72%) Jesus (69%) Lord (68%) Christ (65%) Spirit (55%) Holy (53%) God (50%) Prayer (50%) |
| 37 | Sunday (100%) Tuesday (100%) Thursday (100%) Olympic (100%) Christmas (100%) rugby (100%) Championship (100%) weekends (100%) | days (91%) anniversary (90%) month (88%) week (84%) mpi (83%) semester (81%) mand (80%) Olympics (78%) cent (76%) season (76%) perm (75%) |
| 43 | Armenian (100%) ijan (100%) enia (96%) Iraq (95%) Iranian (92%) Iran (92%) Saudi (90%) northern (90%) Lebanon (90%) Singapore (88%) Turkey (88%)<br>Asia (87%) Egypt (86%) western (86%) | enia (90%) invasion (80%) Arabia (76%) irregular (66%) regions (64%) border (63%) Kong (61%) ians (61%) bases (60%) Republic (59%) Ireland (58%) Korea (58%) War (55%) Carolina (52%) |
| 4 | sq (89%) Main (70%) reversal (69%) YR (63%) GC (56%) Overall (50%) 79 (50%) main (50%) RE (46%) PCR (46%) tomb (45%) normal (43%) intensity (41%) Overall (41%) median (41%) | YR (90%) Character (88%) sq (77%) Os (76%) GHz (71%) fluence (60%) amycin (60%) pixels (56%) = (53%) arc (52%) Story (52%) = (51%) anth (50%) GHz (50%) cm (46%) |
| 0 | ESM (100%) icillin (100%) agra (98%) aust (96%) asa (93%) pills (92%) mg (85%) uk (82%) login (82%) doc (81%) generic (81%) cd (81%) Essay (81%) password (81%) Content (80%) | *, (100%) sil (96%) pills (91%) vi (90%) xen (87%) pharmacy (87%) gener (85%) aust (82%) mg (75%) Content (75%) uk (73%) THAT (73%) dispens (68%) icillin (68%) generic (66%) |
| 3 | grandmother (92%) brother (91%) Daisy (83%) daughter (78%) mum (75%) father (72%) wife (70%) husband (70%) lady (63%) dad (62%) boy (61%) | hood (36%) mother (35%) inde (31%) boy (29%) girl (28%) married (27%) tri (21%) Gab (20%) died (18%) taught (14%) lived (13%) knew (10%) |
| 48 | compared (42%) !) (41%) Then (41%) ’, (40%) ), (35%) ”, (35%) instead (33%) | except (60%) tennis (41%) Marks (40%) Dunn (33%) tears (30%) Arizona (30%) |
| 23 | …. (58%) Therefore (55%) So (46%) !!! (46%) And (44%) According (41%) .” (41%) !! (40%) ?” (38%) But (38%) | (53%) Republican (50%) Jack (47%) THIS (40%) Democratic (40%) according (39%) So (38%) Step (33%) |

We define vocabulary specialization as the proportion of tokens with a token ID $x$ (also called vocabulary element) that are routed to one particular expert $E_{i}$ out of all experts in that layer:

$$ $\text{Vocabulary specialization}(E_{i},x)=\frac{N_{x,E_{i}}^{(k)}}{N_{x}},$ (8) $$

where:

- •
$E_{i}$: The $i$th expert in the model.
- •
$x$: The token ID being analyzed.
- •
$k$: The number of experts considered (e.g., $k=8$ means considering the top 8 experts with the highest routing probabilities).
- •
$N_{x,E_{i}}$: The number of times input data is routed to $E_{i}$ for $x$.
- •
$N_{x}$: The total number of times input data is routed across all experts for $x$.

Vocabulary specialization thus refers to how specialized a particular expert is on some vocabulary item. We distinguish input and output variants of this specialization, where $x$ is either the input token ID or the next output token ID (either the ground-truth next token ID or the token ID predicted by the model). A value of 100% indicates that for all occurrences of that vocabulary element, input data is routed to $E_{i}$, whereas 0% indicates an expert that is fully irrelevant for that vocabulary element and can be effectively removed from the model without affecting performance whenever the token ID appears.

In Figure 23 we find that vocabulary specialization is higher in later layers, similar to how later layers saturate earlier (§5.1). Later layers also specialize more on predicted output token IDs rather than input token IDs, i.e., the routing is decided more by the token the model is about to predict rather than the original input token. This is intuitive as in earlier layers there is more uncertainty about which token the model will predict. At 90%, expert 27 specializes the most, which we find in Table 8 to activate for many non-alphabetic tokens, such as Cyrillic and Devanagari letters. Expert 43 shows specialization on geographic terms in both input and output tokens. Experts 48 and 23 both focus on connector words, such as Then and Therefore. This is likely because they commonly process tokens together with a high co-activation of 60% in Figure 21 (middle). Based on our findings in §5.3 that for GitHub and arXiv often the same experts in layer 7 activate, we display one such expert (expert ID 4) in Table 8. It seems to specialize in measurements, such as sq, YR (year), and GHz. These are common terms in scientific papers corresponding to the arXiv domain and likely also in GitHub code for computations related to measurements. They are less likely to appear in books, which explains the low activation of expert ID 4 in layer 7 for book data in Figure 22. Expert 3 is among the three most active experts of layer 7 for book data in Figure 22 (fourth yellow bar for layer 7). This resonates when looking at its specialization on family terms in Table 8, which are far more common in books than scientific papers or code. Overall, domain specialization and vocabulary specialization are closely linked to one another, as domains are usually characterized by their distinct word distribution. In Appendix G (Figure 32), we link them more closely by comparing the extent of vocabulary specialization across domains and expert IDs. In Appendix G (Figure 30, Figure 31) we also find that OLMoE-1B-7B exhibits stronger vocabulary specialization than Mixtral-8x7B.

## 6 Related Work

##### Advances in MoEs

Current LMs still largely follow the transformer architecture [] with only few architectural changes that have been widely adopted, such as decoder-only training [], SwiGLU activations [, ], RoPE [], MQA/GQA [, ] and RMSNorm []. Model sparsity via Mixture-of-Experts is one modification still under active exploration with some early adoption but most LMs, including Llama 3 [], still rely on a dense architecture. There has been a lot of progress in improving the sparsely-gated MoE layer since its introduction []: New routing techniques [, , , , , , , , ], fine-grained expert segmentation [, ], stability [] and efficiency [, , , , , , , ] improvements. In this work, we perform many experiments to provide insights into training Mixture-of-Experts LMs. Subsequently, we train OLMoE-1B-7B for 5T tokens. No prior MoE has been overtrained [] to this extent to our knowledge making OLMoE-1B-7B the best testbed to research performance saturation of MoEs vs. dense models. With OLMoE we hope to facilitate such and other research to help the field uncover whether MoEs should make it into all future LMs and with what precise configuration.

##### Open LMs

A variety of model families have been proposed under varying degrees of openness commonly categorized based on whether model weights are available. Closed-weight models include GPT [, ], Gemini [, ], PaLM [, ], Reka [], and open-weight ones include Llama [, , ], Mistral [, ], Gemma [, ], Falcon [, ], MPT [], Qwen [, ], GLM [], Yi [], DeepSeek [, , ], Nemotron [, , ], InternLM [], Baichuan [], Phi [, , ], StableLM [], OPT []. However, besides model weights, training data and code are key to enabling scientific research of these models [, ] and distributing their benefits broadly []. There have been few releases also including data and code in addition to model weights which we refer to as “fully open-source”: BLOOM [, , , ], GPT-NeoX [, , ], StarCoder [, , , , ], Pythia [], OLMo [], LLM360 [], Cerebras-GPT [], DCLM [], MAP-Neo [], RWKV [, ], and SmolLM []. For Mixture-of-Experts only OpenMoE [] aims to be fully open-source, however, its poor performance limits its usefulness. We release OLMoE-1B-7B as the first state-of-the-art Mixture-of-Experts LM that is fully open-source: model weights, data, code, and logs.

## 7 Conclusion

We open-source OLMoE-1B-7B and OLMoE-1B-7B-Instruct including model, data, code, and logs. At 1B active and 7B total parameters, our models yield state-of-the-art performance among models with a similar amount of active parameters even outperforming larger models including DeepSeekMoE-16B and Llama2-13B-Chat. We share various training experiments and define and analyze router saturation, expert co-activation, domain and vocabulary specialization of our model. Through our fully open release, we seek to help the field build better MoEs. We are excited about new iterations of OLMoE to close the gap between frontier models and fully open models.

## Author Contributions

Niklas Muennighoff proposed and led the project. He ran the pretraining experiments, pretrained the model, helped run adaptation and analysis, and wrote most of the paper.
Luca Soldaini created the pretraining dataset and advised on pretraining.
Dirk Groeneveld advised on pretraining, especially stability and throughput improvements.
Kyle Lo helped with pretraining dataset creation, analyzed data experiments, and advised on data and framing, and helped edit the paper.
Jacob Morrison co-created the adaptation dataset, ran most adaptation experiments, and helped edit the paper.
Sewon Min analyzed router saturation, expert correlation, and vocabulary specialization, and helped frame and edit the paper.
Weijia Shi analyzed domain and vocabulary specialization, advised at various project stages, and helped edit the paper.
Pete Walsh advised on pretraining, especially stability and throughput improvements.
Oyvind Tafjord ran OLMES evaluations.
Nathan Lambert co-created the adaptation dataset, advised on adaptation, and helped edit the paper.
Yuling Gu ran OLMES evaluations and helped edit the paper.
Shane Arora uploaded the models and helped with code review.
Akshita Bhagia supported stability investigations and helped with DCLM evaluations.
Dustin Schwenk supported stability investigations.
David Wadden ran DCLM evaluations and helped with Weights & Biases reports.
Alexander Wettig analyzed load balancing, routing, and domain specialization, and helped edit the paper.
Binyuan Hui advised on pretraining.
Tim Dettmers advised on analysis and inference experiments.
Douwe Kiela advised on framing.
Ali Farhadi advised on pretraining and framing.
Noah A. Smith advised on pretraining, and helped frame and edit the paper.
Pang Wei Koh advised on analysis, and helped frame and edit the paper.
Amanpreet Singh advised on pretraining, framing and helped edit the paper.
Hannaneh Hajishirzi was responsible for direction and advising of the overall effort and helped frame and edit the paper.

## Acknowledgements

OLMoE would not be possible without the support of many individuals and institutions. We thank our teammates at the Allen Institute for AI, Contextual AI, and the University of Washington for their support, especially Aditya Kusupati, Ananya Harsh Jha, Caitlin Wittlif, Carissa Schoenick, Costa Huang, Crystal Nam, David Atkinson, Emma Strubell, Faeze Brahman, Hamish Ivison, Karel D’Oosterlinck, Matt Latzke, Ian Magnusson, Jack Merullo, Jay Chen, Jennifer Dumas, Jiacheng Liu, Johann Dahm, Luke Zettlemoyer, Michael Schmitz, Michael Wilson, Pradeep Dasigi, Sahil Verma, Sam Skjonsberg, Sophie Lebrecht, Stas Bekman, Taira Anderson, Valentina Pyatkin, Yanai Elazar, Yizhong Wang, and Yoganand Chandrasekhar. We also thank Armen Aghajanyan, Akshat Shrivastava, Colin Raffel, Haokun Liu, Ludwig Schmidt, and Shayne Longpre. PWK is supported by the Singapore National Research Foundation and the National AI Group in the Singapore Ministry of Digital Development and Innovation under the AI Visiting Professorship Programme (award number AIVP-2024-001).

## Appendix A Artifacts

**Table 9: All artifacts released and used in this work. We point from the name used for a given artifact in this work (e.g. Figure 1) to the URL where it can be obtained.**
| Artifact | Public link |
| --- | --- |
| OLMoE-1B-7B | [https://hf.co/allenai/OLMoE-1B-7B-0924](https://hf.co/allenai/OLMoE-1B-7B-0924) |
| OLMoE-1B-7B-Instruct | [https://hf.co/allenai/OLMoE-1B-7B-0924-Instruct](https://hf.co/allenai/OLMoE-1B-7B-0924-Instruct) |
| OLMoE-1B-7B-SFT | [https://hf.co/allenai/OLMoE-1B-7B-0924-SFT](https://hf.co/allenai/OLMoE-1B-7B-0924-SFT) |
| OLMoE-Mix | [https://hf.co/datasets/allenai/OLMoE-mix-0924](https://hf.co/datasets/allenai/OLMoE-mix-0924) |
| SFT data | [https://hf.co/datasets/allenai/](https://hf.co/datasets/allenai/tulu-v3.1-mix-preview-4096-OLMoE) |
| [tulu-v3.1-mix-preview-4096-OLMoE](https://hf.co/datasets/allenai/tulu-v3.1-mix-preview-4096-OLMoE) |  |
| KTO/DPO data | [https://hf.co/datasets/allenai/](https://hf.co/datasets/allenai/ultrafeedback_binarized_cleaned) |
| [ultrafeedback_binarized_cleaned](https://hf.co/datasets/allenai/ultrafeedback_binarized_cleaned) |  |
| Code | [https://github.com/allenai/OLMoE](https://github.com/allenai/OLMoE) |
| Logs | [https://wandb.ai/ai2-llm/olmoe/reports/](https://wandb.ai/ai2-llm/olmoe/reports/OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3) |
| [OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3](https://wandb.ai/ai2-llm/olmoe/reports/OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3) |  |
| BLOOM-7B | [https://hf.co/bigscience/bloom-7b1](https://hf.co/bigscience/bloom-7b1) |
| DeepSeekMoE-3B-16B | [https://hf.co/deepseek-ai/deepseek-moe-16b-base](https://hf.co/deepseek-ai/deepseek-moe-16b-base) |
| DeepSeekMoE-3B-16B+chat | [https://hf.co/deepseek-ai/deepseek-moe-16b-chat](https://hf.co/deepseek-ai/deepseek-moe-16b-chat) |
| DCLM-1B | [https://hf.co/TRI-ML/DCLM-1B](https://hf.co/TRI-ML/DCLM-1B) |
| DCLM-7B | [https://hf.co/TRI-ML/DCLM-7B](https://hf.co/TRI-ML/DCLM-7B) |
| Falcon-7B | [https://hf.co/tiiuae/falcon-7b](https://hf.co/tiiuae/falcon-7b) |
| Gemma2-3B | [https://hf.co/google/gemma-2-2b](https://hf.co/google/gemma-2-2b) |
| Gemma2-9B | [https://hf.co/google/gemma-2-9b](https://hf.co/google/gemma-2-9b) |
| JetMoE-2B-9B | [https://hf.co/jetmoe/jetmoe-8b](https://hf.co/jetmoe/jetmoe-8b) |
| JetMoE-2B-9B+SFT | [https://hf.co/jetmoe/jetmoe-8b-sft](https://hf.co/jetmoe/jetmoe-8b-sft) |
| JetMoE-2B-9B+Chat | [https://hf.co/jetmoe/jetmoe-8b-chat](https://hf.co/jetmoe/jetmoe-8b-chat) |
| Llama-7B | [https://hf.co/huggyllama/llama-7b](https://hf.co/huggyllama/llama-7b) |
| Llama2-7B | [https://hf.co/meta-llama/Llama-2-7b-hf](https://hf.co/meta-llama/Llama-2-7b-hf) |
| Llama3.1-8B | [https://hf.co/meta-llama/Meta-Llama-3.1-8B](https://hf.co/meta-llama/Meta-Llama-3.1-8B) |
| MPT-7B | [https://hf.co/mosaicml/mpt-7b](https://hf.co/mosaicml/mpt-7b) |
| Mistral-7B | [https://hf.co/mistralai/Mistral-7B-v0.1](https://hf.co/mistralai/Mistral-7B-v0.1) |
| Mixtral-8x7B | [https://hf.co/mistralai/Mixtral-8x7B-v0.1](https://hf.co/mistralai/Mixtral-8x7B-v0.1) |
| OLMo-1B (0724) | [https://hf.co/allenai/OLMo-1B-0724-hf](https://hf.co/allenai/OLMo-1B-0724-hf) |
| OLMo-7B (0724) | [https://hf.co/allenai/OLMo-7B-0724-hf](https://hf.co/allenai/OLMo-7B-0724-hf) |
| OpenMoE-3B-9B | [https://hf.co/OrionZheng/openmoe-8b](https://hf.co/OrionZheng/openmoe-8b) |
| Pythia-7B | [https://hf.co/EleutherAI/pythia-6.9b](https://hf.co/EleutherAI/pythia-6.9b) |
| Qwen1.5-3B-14B | [https://hf.co/Qwen/Qwen1.5-MoE-A2.7B](https://hf.co/Qwen/Qwen1.5-MoE-A2.7B) |
| Qwen1.5-3B-14B+Chat | [https://hf.co/Qwen/Qwen1.5-MoE-A2.7B-Chat](https://hf.co/Qwen/Qwen1.5-MoE-A2.7B-Chat) |
| StableLM2-2B | [https://hf.co/stabilityai/stablelm-2-1_6b](https://hf.co/stabilityai/stablelm-2-1_6b) |
| TinyLlama-1B | [https://hf.co/TinyLlama/TinyLlama_v1.1](https://hf.co/TinyLlama/TinyLlama_v1.1) |

## Appendix B Training Configuration

##### Pretraining

We display the pretraining hyperparameter configuration of OLMoE-1B-7B in Appendix B comparing with other relevant models. We follow Groeneveld et al. [] using the AdamW optimizer [] with ZeRO [] via PyTorch FSDP [] and mixed-precision training []. Our main model settings differing from Groeneveld et al. [] are: (1) MoE-related changes: OLMoE-1B-7B is a sparsely activated decoder-only transformer [] using dropless Mixture-of-Experts []. Unlike most prior MoEs, we use a high granularity [, ] with 64 small experts with an FFN dimension of just 1,024 rather than a few large experts. We further use two auxiliary losses: router z-loss [] and load balancing loss []. (2) Stability improvements: (a) We use a truncated normal initialization with a standard deviation of 0.02 and a minimum (maximum) cut-off of -0.06 (0.06) corresponding to three standard deviations. (b) We use QK normalization [, , ].
(c) We use RMSNorm [] instead of the non-parametric LayerNorm used in Groeneveld et al. []. (3) Performance improvements: Besides some of the stability improvements which also impact performance, we also reduce the AdamW epsilon to 1.0E-08 from the 1.0E-05 used in Groeneveld et al. [] to speed up convergence. Finally, we train OLMoE-1B-7B for significantly longer than all prior OLMo models amounting to 5T tokens and thus more than one epoch (1.2) following Muennighoff et al. []. We shuffle the pretraining dataset before starting the second epoch. For the final 100B tokens, we decay the learning rate linearly from 5.0E-04 to 0 (“annealing”). We experiment with many of these settings in §4.

##### Adaptation

For finetuning we use Open Instruct [, ].(^7^77Code: [https://github.com/allenai/open-instruct](https://github.com/allenai/open-instruct)) We filter all SFT samples to a length of fewer than 4096 tokens to match the sequence length of the model. Following Muennighoff et al. [], we aggregate loss at the token level during SFT to improve performance on long generative tasks, such as AlpacaEval. We finetune in BF16 with a global batch size of 128 (4 H100 nodes with 8 GPUs each, a per device batch size of 2, and 2 gradient accumulation steps). We train for 2 epochs with a constant learning rate of 2.0E-5. For DPO [], we reduce the global batch size to 32 (4 H100 nodes with 8 GPUs each and a per device batch size of 1). We train for 3 epochs with a learning rate of 5.0E-7 and a DPO beta of 0.1. Our adapted models are built on top of our annealed checkpoint, and we include the load balancing loss during both SFT and DPO based on our experiments in §4.3. Our preference tuning recipe is heavily optimized for DPO based on extensive experiments by Ivison et al. [], thus for KTO [] we experiment with a few settings in Appendix F. Our final KTO adaptation uses the same hyperparameters as DPO, except that we use the RMSProp optimizer instead of Adam, which we use for SFT and DPO, and that we reduce the training duration to 1.3 epochs (5,000 steps) for KTO instead of the 3 epochs used for DPO.

##### Hardware

We pretrain OLMoE-1B-7B on 256 H100 GPUs for approximately 10 days with NV-link interconnect across GPUs and InfiniBand interconnect across nodes. We also use H100 GPUs for all our experiments but some use a cluster with GCP TCPx interconnect across nodes instead. For adaptation, we use 32 H100 GPUs for 33 hours to instruction tune and for another 14 hours to preference tune via DPO. For KTO adaptation we use 8 H100 GPUs for 30 hours instead.

**Table 10: Pretraining hyperparameters of OLMoE-1B-7B and comparable models trained from scratch. We highlight rows where OLMoE-1B-7B differs from OLMo-1B. Active params include vocab params. “?” = undisclosed settings, FFN = feed-forward network, Attn = Attention, LR = learning rate, WSD = Weight-Stable-Decay [], LBL = load balancing loss, Inv Sq Root = Inverse Square Root decay [], trunc = truncation, std = standard deviation, “varies” = stds that are layer or weight-dependent.**
|  | OLMoE-1B-7B | JetMoE | OpenMoE | OLMo-1B (0724) |
| --- | --- | --- | --- | --- |
| Dimension | 2,048 | 2,048 | 2,048 | 2,048 |
| Activation | SwiGLU | SwiGLU | SwiGLU | SwiGLU |
| FFN dimension | 1,024 | 5,632 | 8,192 | 8,192 |
| Vocab size | 50,304 | 32,000 | 256,384 | 50,304 |
| Attn heads | 16 | 16 | 24 | 16 |
| Num layers | 16 | 24 | 32 | 16 |
| Layer norm type | RMSNorm | RMSNorm | RMSNorm | non-parametric |
| Layer norm eps | 1.0E-05 | 1.0E-05 | 1.0E-06 | 1.0E-05 |
| QK-Norm | yes | no | no | no |
| Pos emb. | RoPE | RoPE | RoPE | RoPE |
| RoPE $\theta$ | 10,000 | 10,000 | 10,000 | 10,000 |
| Attention variant | full | MoA | full | full |
| Biases | - | MLP & Attn | - | - |
| Weight tying | no | yes | no | no |
| Init dist | trunc normal | ? | ? | normal |
| Init std | 0.02 | 0.02 | varies | varies |
| Init trunc | 3std | - | - | - |
| MoE layers | Every | Every | Every 6th | - |
| MoE layer type | dMoE | dMoE | ST-MoE | - |
| # Experts | 64 | 8 | 32 | 1 |
| # Activated | 8 | 2 | 2 | 1 |
| # Vocab params | 103M | 66M | 525M | 103M |
| # Active params | 1.3B | 2.2B | 2.6B | 1.3B |
| # Total params | 6.9B | 8.5B | 8.7B | 1.3B |
| Sequence length | 4,096 | 4,096 | 2,048 | 4,096 |
| Batch size (samples) | 1,024 | 1,024 | 2,048 | 512 |
| Batch size (tokens) | 4M | 4M | 4M | 2M |
| warmup steps | 2,000 | 2,500 | 10,000 | 2,000 |
| peak LR | 4.0E-04 | 5.0E-04 | 0.01 | 4.0E-04 |
| minimum LR | 5.0E-04 | 5.0E-05 | - | 5.0E-05 |
| optimizer | AdamW | AdamW | Adafactor | AdamW |
| weight decay | 0.1 | 0.1 | 0.0 | 0.1 |
| beta1 | 0.9 | ? | 0.9 | 0.9 |
| beta2 | 0.95 | ? | - | 0.95 |
| AdamW epsilon | 1.0E-08 | ? | - | 1.0E-05 |
| LR schedule | cosine | WSD | Inv Sq Root | cosine |
| gradient clipping | global 1.0 | global 1.0 | global 1.0 | global 1.0 |
| gradient reduce dtype | FP32 | ? | ? | FP32 |
| optimizer state dtype | FP32 | ? | ? | FP32 |
| LBL weight | 0.01 | 0.01 | 0.01 | - |
| Router z-loss weight | 0.001 | 0.001 | 0.0001 | - |
| Pretraining tokens | 5,033B | 1,000B | 1,100B | 2,000B |
| Annealing tokens | 100B | 250B | - | 50B |
| Annealing schedule | linear | - | - | linear |
| Annealing min LR | 0 | - | - | 0 |

## Appendix C Evaluation Setup

**Table 11: Summary of downstream evaluation during and after pretraining (OLMES). ARC-C and ARC-E refer to ARC-Challenge and -Easy, CSQA=CommonsenseQA, OBQA=OpenBookQA, CF=Completion/Cloze formulation, MCF=Multiple-choice formulation, pmi=pointwise-mutual-information, Var=variants referring to the use of few-shots varying from 0-5.**
| Dataset ($\downarrow$) | During pretraining | After pretraining (OLMES []) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Format | Shot | Norm | Split | Format | Shot | CF Norm | Split |  |
| ARC-C [] | CF | 0 | token | val | max(MCF,CF) | 5 | pmi | test |
| ARC-E [] | CF | 0 | none | val | max(MCF,CF) | 5 | character | test |
| BoolQ [] | CF | 0 | none | val | max(MCF,CF) | 5 | none | val |
| COPA [] | CF | 0 | none | val | - | - | - | - |
| CSQA [] | CF | 0 | token | val | max(MCF,CF) | 5 | pmi | val |
| HellaSwag [] | CF | 0 | token | val | max(MCF,CF) | 5 | character | val |
| MMLU [] | MCF | 5 | none | val | max(MCF,CF) | 5 | character | test |
| MMLU Var | CF | 0-5 | token | val | - | - | - | - |
| OBQA [] | CF | 0 | token | val | max(MCF,CF) | 5 | pmi | test |
| PIQA [] | CF | 0 | token | val | max(MCF,CF) | 5 | character | val |
| SciQ [] | CF | 0 | none | val | - | - | - | - |
| SocialIQA [] | CF | 0 | token | val | max(MCF,CF) | 5 | character | val |
| Winogrande [] | CF | 0 | none | val | max(MCF,CF) | 5 | none | val |

##### During pretraining

We evaluate using a similar in-loop evaluation setup as Groeneveld et al. [], with the addition of more tasks such as CommonsenseQA, PIQA, and different implementations of MMLU. Following Groeneveld et al. [], for the majority of the tasks, we perform 0-shot evaluation using the Completion/Cloze formulation (CF), ranking each answer string using language model probabilities. In terms of probability normalization, there is either no normalization (none) or normalization by the number of tokens in the answer (token) when ranking solely based on probability may heavily favor shorter answers []. For MMLU, the in-loop evaluation also includes a setup where we increase the total number of instances by including a range of 0-shot to 5-shot setups together as we found this provides smoother trends as the training proceeds (“MMLU Var”). We also include the Multiple-choice formulation (MCF) version of MMLU, scoring prediction of answer labels like A/B/C/D, which generally starts to rise only later in training as models only gain the multiple-choice capability later (at around 1T tokens for OLMoE-1B-7B in Figure 25). We also evaluate perplexity on selected validation sets from Paloma [, , , , , ]. All code used for evaluation during pretraining is at [https://github.com/allenai/OLMo/tree/61ac104d616ec5435db225796e5c7532c9abd95a/olmo/eval](https://github.com/allenai/OLMo/tree/61ac104d616ec5435db225796e5c7532c9abd95a/olmo/eval).

##### After pretraining - OLMES

We perform evaluations following the OLMES evaluation standard [], with the suite of tasks in the original paper. OLMES (Open Language Model Evaluation Standard) is a standard for reproducible LM evaluations that is open, practical, and documented, providing recommendations guided by experiments and results from the literature [, , ]. It is designed to support comparisons between smaller base models that require the Cloze formulation of multiple-choice questions against larger models that can utilize the Multiple-choice formulation. To make our evaluations reproducible, we follow OLMES in prompt formatting, choice of in-context examples, probability normalization, task formulation, as well as all other details. We summarize this setup in Table 4 and refer to Gu et al. [] for more details.

##### After pretraining - DCLM

For results on the DCLM tasks [] in Table 13, we precisely follow their setup using the evaluation code released by the authors at [https://github.com/mlfoundations/dclm](https://github.com/mlfoundations/dclm). “Core” results are the low variance tasks in their evaluation code, while “Extended” corresponds to the heavy tasks.

##### After adaptation

After supervised finetuning and direct preference optimization, we evaluate models using a subset of the evaluations and the same overall setup used in Ivison et al. [] and Wang et al. []. We cover a wide range of model capabilities in our evaluation suite including coding (HumanEval []), general and mathematical reasoning (Big Bench Hard [], GSM8k []), world knowledge (MMLU), general instruction following (AlpacaEval 1.0 [], not the length-controlled variant []), precise instruction following (IFEval []) and safety (XSTest []). We refer to Wang et al. [] for more details on each benchmark.

## Appendix D Openness of Models

We list the openness of various models summarized in Figure 1. We exclude Switch Transformers [], as it was published over three years ago and is very different from more recent MoE models (MLM objective, Encoder-decoder, etc.).

##### Grok-86B-314B [ 195 ]

- •
Model: Their model is licensed under the open-source Apache 2.0 license.
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### Mixtral-39B-141B and Mixtral-13B-42B [ 78 ]

- •
Model: Their model is licensed under the open-source Apache 2.0 license.
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### DBRX-36B-132B [ 40 ]

- •
Model: The model is licensed under a custom non-open-source license(^8^88[https://www.databricks.com/legal/open-model-license](https://www.databricks.com/legal/open-model-license)) with additional use-case restrictions.(^9^99[https://www.databricks.com/legal/acceptable-use-policy-open-model](https://www.databricks.com/legal/acceptable-use-policy-open-model))
- •
Data: Unavailable.
- •
Code: They use closed-source custom adaptations of their public libraries LLM-foundry, composer, and megablocks.(^10^1010[https://github.com/databricks/dbrx](https://github.com/databricks/dbrx))
- •
Logs: Unavailable.

##### Skywork-MoE-22B-146B [ 190 ]

- •
Model: The model is licensed under a custom non-open-source license.(^11^1111[https://github.com/SkyworkAI/Skywork/blob/main/Skywork%20Community%20License.pdf](https://github.com/SkyworkAI/Skywork/blob/main/Skywork%20Community%20License.pdf))
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### DeepSeekV2-21B-236B [ 43 ] and DeepSeekMoE-3B-14B [ 39 ]

- •
Model: The models are licensed under custom non-open-source licenses.(^12^1212[https://github.com/deepseek-ai/DeepSeek-MoE/blob/main/LICENSE-MODEL](https://github.com/deepseek-ai/DeepSeek-MoE/blob/main/LICENSE-MODEL) and [https://github.com/deepseek-ai/DeepSeek-V2/blob/main/LICENSE-MODEL](https://github.com/deepseek-ai/DeepSeek-V2/blob/main/LICENSE-MODEL))
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### Arctic-17B-480B [ 160 ]

- •
Model: The model is licensed under the open-source Apache 2.0 license.
- •
Data: They describe their mixture but do not release it.(^13^1313[https://medium.com/snowflake/snowflake-arctic-cookbook-series-arctics-approach-to-data-b81a8a0958bd](https://medium.com/snowflake/snowflake-arctic-cookbook-series-arctics-approach-to-data-b81a8a0958bd))
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### Qwen2-14B-57B [ 178 ]

- •
Model: The model is licensed under the open-source Apache 2.0 license.
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### Jamba-12B-52B [ 96 ]

- •
Model: The model is licensed under the open-source Apache 2.0 license.
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### Qwen1.5-3B-14B [ 178 ]

- •
Model: The model is licensed under a custom non-open-source license.(^14^1414[https://hf.co/Qwen/Qwen1.5-MoE-A2.7B/blob/main/LICENSE](https://hf.co/Qwen/Qwen1.5-MoE-A2.7B/blob/main/LICENSE))
- •
Data: Unavailable.
- •
Code: Unavailable.
- •
Logs: Unavailable.

##### JetMoE-2B-9B [ 156 ]

- •
Model: The model is licensed under the open-source Apache 2.0 license.
- •
Data: They describe their mixture but do not release it.
- •
Code: They make their fork of megablocks publicly available,(^15^1515[https://github.com/yikangshen/megablocks](https://github.com/yikangshen/megablocks)) however, their Megatron-LM training code is not available.(^16^1616[https://hf.co/jetmoe/jetmoe-8b/discussions/5#661ee52c03251697a0b155cc](https://hf.co/jetmoe/jetmoe-8b/discussions/5#661ee52c03251697a0b155cc))
- •
Logs: Unavailable.

##### OpenMoE-2B-9B [ 198 ]

- •
Model: The model is licensed under the open-source Apache 2.0 license.
- •
Data: They make scripts for recreating their data available.
- •
Code: They make their code available.(^17^1717[https://github.com/XueFuzhao/OpenMoE/tree/main?tab=readme-ov-file#training-with-tpugpu](https://github.com/XueFuzhao/OpenMoE/tree/main?tab=readme-ov-file#training-with-tpugpu))
- •
Logs: Unavailable.

##### OLMoE-1B-7B

- •
Model: The model is licensed under the open-source Apache 2.0 license.
- •
Data: The data is licensed under the open-source ODC-By 1.0 license.
- •
Code: The code is licensed under the open-source Apache 2.0 license.
- •
Logs: Logs are available with the same open-source license as the code (Apache 2.0).

## Appendix E Additional Evaluation

Figure: Figure 24: Losses of OLMoE-1B-7B during training. The Books, Reddit, and Stack [] datasets are from Dolma 1.7 [] via Paloma []. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B--Vmlldzo4OTcyMjU3](https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B--Vmlldzo4OTcyMjU3)
Refer to caption: x29.png

Figure: Figure 25: Evaluation of OLMoE-1B-7B and the current best OLMo models during pretraining. Grey vertical lines correspond to where the respective run enters annealing with the 1st line being for OLMo-7B, the 2nd for OLMo-1B, and the third for OLMoE-1B-7B. Figure 3 is a version of this plot with training FLOPs as the x-axis. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz](https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz)
Refer to caption: x30.png

**Table 12: More results on OLMES. ${}^{\text{\textdagger}}$ indicates use of the MCF score, see Appendix C.**
| Model | ARC_C | ARC_E | BoolQ | CSQA | HSwag | MMLU | OBQA | PIQA | SIQA | WinoG | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LMs with 7-9B active parameters |  |  |  |  |  |  |  |  |  |  |  |
| Mistral-7B | $78.6^{\text{\textdagger}}$ | $90.8^{\text{\textdagger}}$ | 89.3 | $72.4^{\text{\textdagger}}$ | 83.0 | $64.0^{\text{\textdagger}}$ | $80.6^{\text{\textdagger}}$ | 82.8 | $71.3^{\text{\textdagger}}$ | 77.9 | 79.1 |
| OLMo-7B (0724) | $68.0^{\text{\textdagger}}$ | $85.7^{\text{\textdagger}}$ | 85.3 | $85.4^{\text{\textdagger}}$ | 80.5 | $54.9^{\text{\textdagger}}$ | $67.6^{\text{\textdagger}}$ | 79.3 | $76.1^{\text{\textdagger}}$ | 73.2 | 75.6 |
| DCLM-7B | $79.8^{\text{\textdagger}}$ | $92.3^{\text{\textdagger}}$ | 87.0 | 77.0 | 82.3 | $64.4^{\text{\textdagger}}$ | $79.6^{\text{\textdagger}}$ | 80.1 | $71.2^{\text{\textdagger}}$ | 77.3 | 79.1 |
| Llama2-7B | 54.2 | 84.0 | 86.1 | 74.2 | 78.9 | $46.2^{\text{\textdagger}}$ | 57.8 | 77.5 | 59.6 | 71.7 | 69.0 |
| Llama3.1-8B | $79.5^{\text{\textdagger}}$ | $91.7^{\text{\textdagger}}$ | 88.5 | $74.3^{\text{\textdagger}}$ | 81.6 | $66.9^{\text{\textdagger}}$ | $78.6^{\text{\textdagger}}$ | 81.1 | $71.4^{\text{\textdagger}}$ | 76.6 | 79.0 |
| Gemma2-9B | $89.5^{\text{\textdagger}}$ | $95.5^{\text{\textdagger}}$ | 89.4 | $78.8^{\text{\textdagger}}$ | $87.3^{\text{\textdagger}}$ | $70.6^{\text{\textdagger}}$ | $88.4^{\text{\textdagger}}$ | $86.1^{\text{\textdagger}}$ | $76.0^{\text{\textdagger}}$ | 78.8 | 84.0 |
| LMs with 2-3B active parameters |  |  |  |  |  |  |  |  |  |  |  |
| StableLM-2B | $50.6^{\text{\textdagger}}$ | 75.3 | 82.3 | $70.4^{\text{\textdagger}}$ | 70.3 | $40.4^{\text{\textdagger}}$ | $56.6^{\text{\textdagger}}$ | 75.6 | $64.3^{\text{\textdagger}}$ | 65.8 | 65.1 |
| Gemma2-3B | $67.5^{\text{\textdagger}}$ | $84.3^{\text{\textdagger}}$ | 83.6 | $66.4^{\text{\textdagger}}$ | 74.6 | $53.3^{\text{\textdagger}}$ | $68.8^{\text{\textdagger}}$ | 78.5 | $64.7^{\text{\textdagger}}$ | 71.8 | 71.4 |
| JetMoE-2B-9B | $61.4^{\text{\textdagger}}$ | $81.9^{\text{\textdagger}}$ | 85.7 | $75.3^{\text{\textdagger}}$ | 81.7 | $49.1^{\text{\textdagger}}$ | $68.0^{\text{\textdagger}}$ | 80.3 | $71.3^{\text{\textdagger}}$ | 70.7 | 72.5 |
| OpenMoE-3B-9B | 29.3 | 50.6 | 63.2 | 21.5 | 44.4 | 27.4 | 34.6 | 63.3 | 42.9 | $51.9^{\text{\textdagger}}$ | 42.9 |
| DeepSeek-3B-16B | 53.4 | 82.7 | 81.9 | 72.7 | 80.4 | $45.5^{\text{\textdagger}}$ | 58.4 | 80.1 | 59.9 | 73.2 | 68.8 |
| Qwen1.5-3B-14B | $77.4^{\text{\textdagger}}$ | $91.6^{\text{\textdagger}}$ | 85.0 | $81.4^{\text{\textdagger}}$ | 80.0 | $62.4^{\text{\textdagger}}$ | $80.6^{\text{\textdagger}}$ | 81.0 | $74.1^{\text{\textdagger}}$ | 72.3 | 78.6 |
| LMs with 1B active parameters |  |  |  |  |  |  |  |  |  |  |  |
| OLMo-1B (0724) | 36.4 | 53.5 | 66.8 | 42.4 | 67.5 | 32.1 | 44.2 | 74.0 | 45.2 | 62.9 | 52.5 |
| TinyLlama-1B | 38.1 | 69.5 | 63.6 | 61.1 | 60.8 | 33.6 | 45.0 | 71.7 | 50.4 | 60.1 | 55.4 |
| Pythia-1B | 31.4 | 63.4 | $56.8^{\text{\textdagger}}$ | 50.9 | 48.0 | 31.1 | 40.4 | 68.9 | 46.4 | 52.7 | 49.0 |
| DCLM-1B | $57.6^{\text{\textdagger}}$ | 79.5 | 80.9 | 71.3 | 75.1 | $48.5^{\text{\textdagger}}$ | $60.0^{\text{\textdagger}}$ | 76.6 | $60.5^{\text{\textdagger}}$ | 68.1 | 67.8 |
| OLMoE-1B-7B | $62.1^{\text{\textdagger}}$ | 84.2 | 79.2 | 72.9 | 80.0 | $54.1^{\text{\textdagger}}$ | $65.4^{\text{\textdagger}}$ | 79.8 | $63.0^{\text{\textdagger}}$ | 70.2 | 71.1 |

**Table 13: DCLM evaluation metrics on the Core and Extended task subsets []. ^∗=Core tasks. “annealed” is the final pretraining checkpoint we use for OLMoE-1B-7B and was annealed from the checkpoint at step 1,200,000. We left the non-annealing pretraining run train a little longer resulting in the 1,220,000 checkpoint.**
| OLMoE-1B-7B checkpoint () | step 1,200,000 | step 1,220,000 | annealed | OLMo-1B | OLMo-7B |
| --- | --- | --- | --- | --- | --- |
| AGI Eval LSAT-AR^∗ | 24.3 | 26.5 | 28.7 | 28.3 | 28.3 |
| AGI Eval LSAT-LR | 40.2 | 38.6 | 37.3 | 30.2 | 42.9 |
| AGI Eval LSAT-RC | 47.4 | 43.7 | 46.6 | 23.5 | 61.6 |
| AGI Eval SAT-En | 55.3 | 54.9 | 52.9 | 28.2 | 73.8 |
| AGI Eval SAT-Math CoT | 5.5 | 4.1 | 6.4 | 1.8 | 6.8 |
| AQuA CoT | 2.4 | 2.9 | 2.0 | 2.9 | 6.1 |
| ARC Challenge^∗ | 53.3 | 53.4 | 53.8 | 34.6 | 48.1 |
| ARC Easy^∗ | 77.1 | 78.5 | 77.7 | 64.4 | 75.9 |
| BBQ | 49.8 | 48.3 | 50.6 | 45.8 | 67.2 |
| BigBench CS Algorithms^∗ | 47.1 | 50.2 | 47.2 | 47.5 | 53.6 |
| BigBench Conceptual Combinations | 51.5 | 50.5 | 56.3 | 31.1 | 68.0 |
| BigBench Conlang Translation | 3.7 | 6.1 | 7.3 | 4.3 | 7.3 |
| BigBench Dyck Languages^∗ | 19.3 | 15.9 | 21.5 | 26.6 | 22.2 |
| BigBench Elementary Math QA | 26.2 | 27.0 | 26.9 | 26.2 | 30.4 |
| BigBench Language Identification^∗ | 31.9 | 34.0 | 31.0 | 27.0 | 39.1 |
| BigBench Logical Deduction | 26.6 | 25.3 | 24.6 | 23.6 | 27.3 |
| BigBench Misconceptions | 59.8 | 55.3 | 62.6 | 55.7 | 58.0 |
| BigBench Novel Concepts | 62.5 | 62.5 | 65.6 | 43.8 | 53.1 |
| BigBench Operators^∗ | 36.2 | 34.3 | 33.8 | 23.8 | 45.2 |
| BigBench QA Wikidata^∗ | 68.2 | 68.8 | 69.2 | 67.0 | 69.9 |
| BigBench Repeat Copy Logic^∗ | 15.6 | 15.6 | 18.8 | 3.1 | 9.4 |
| BigBench Strange Stories | 66.7 | 68.4 | 69.5 | 53.4 | 66.1 |
| BigBench Strategy QA | 56.2 | 58.1 | 57.0 | 51.5 | 68.6 |
| BigBench Understanding Fables | 47.1 | 44.4 | 47.6 | 28.0 | 61.4 |
| BoolQ^∗ | 73.3 | 72.8 | 73.2 | 63.7 | 83.9 |
| COPA^∗ | 81.0 | 80.0 | 78.0 | 75.0 | 77.0 |
| CoQA^∗ | 43.7 | 44.4 | 43.7 | 3.4 | 45.4 |
| CommonsenseQA^∗ | 67.2 | 67.0 | 69.3 | 19.6 | 86.0 |
| Enterprise PII Classification | 52.3 | 53.7 | 52.2 | 57.3 | 50.6 |
| GPQA Diamond | 22.2 | 21.2 | 19.7 | 19.7 | 20.2 |
| GPQA Main | 24.8 | 22.3 | 22.5 | 20.3 | 23.0 |
| GSM8K CoT | 6.4 | 7.4 | 7.4 | 4.9 | 30.6 |
| HellaSwag 0-shot^∗ | 76.0 | 76.0 | 77.0 | 65.8 | 76.7 |
| HellaSwag 10-shot^∗ | 77.6 | 77.5 | 78.6 | 66.3 | 78.9 |
| Jeopardy^∗ | 48.8 | 48.7 | 50.3 | 22.6 | 46.5 |
| LAMBADA^∗ | 72.7 | 72.2 | 73.3 | 61.1 | 71.8 |
| LogiQA | 34.9 | 34.3 | 34.6 | 28.7 | 31.0 |
| MMLU Few-shot | 52.2 | 51.9 | 53.3 | 28.4 | 55.1 |
| MMLU Zero-shot | 41.6 | 42.7 | 43.3 | 26.2 | 50.0 |
| Math QA | 26.4 | 27.1 | 27.5 | 24.1 | 29.8 |
| OpenBookQA^∗ | 41.4 | 44.0 | 44.8 | 36.6 | 43.4 |
| PIQA^∗ | 81.3 | 81.2 | 82.0 | 76.4 | 81.7 |
| PubMedQA | 56.1 | 46.6 | 57.9 | 0.2 | 57.9 |
| SQuAD^∗ | 52.9 | 52.4 | 52.4 | 0.0 | 65.5 |
| SVAMP CoT | 30.0 | 28.0 | 33.0 | 14.3 | 44.7 |
| Simple Arithmetic, no spaces | 17.6 | 18.1 | 20.1 | 1.2 | 15.3 |
| Simple Arithmetic, with spaces | 19.5 | 20.6 | 22.1 | 1.8 | 16.0 |
| Social IQA | 71.5 | 70.7 | 69.3 | 69.5 | 84.4 |
| Trivia QA | 54.2 | 53.0 | 55.9 | 25.1 | 51.8 |
| Winogender Female | 50.0 | 46.7 | 50.0 | 41.7 | 58.3 |
| Winogender Male | 55.0 | 58.3 | 60.0 | 63.3 | 58.3 |
| Winograd^∗ | 82.8 | 83.2 | 84.6 | 79.9 | 83.2 |
| Winogrande^∗ | 68.0 | 68.5 | 69.0 | 61.8 | 67.6 |
| Core | 46.3 | 46.5 | 47.2 | 30.2 | 49.8 |
| Extended | 31.3 | 30.9 | 32.5 | 16.9 | 37.0 |

## Appendix F Additional Experiments

Figure: Figure 26: Adding Reddit or FLAN to OLMoE-Mix. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Adding-Reddit-FLAN--Vmlldzo4OTg1NTg4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Adding-Reddit-FLAN--Vmlldzo4OTg1NTg4)
Refer to caption: x31.png

##### Adding Reddit or FLAN to OLMoE-Mix

In Figure 26 we benchmark adding the Reddit or FLAN [] subsets of Dolma 1.7 [] to our pretraining data mix (§2). Overall, we do not find either one to lead to consistent gains, thus we do not use them in our final data mix.

Figure: Figure 27: Load balancing precision. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-FP32-LBL--Vmlldzo4NDMxNDA4](https://wandb.ai/ai2-llm/olmoe/reports/Plot-FP32-LBL--Vmlldzo4NDMxNDA4)
Refer to caption: x32.png

##### Load balancing precision

Fedus et al. [] selectively perform operations related to routing in full precision (FP32) to improve stability. In Figure 27, we test whether computing the load balancing loss in full precision improves stability, but do not find it to reduce spikes. Thus, we stick with bfloat16 (BF16).

Figure: Figure 28: Adding noise to the upcycled checkpoint. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Noise-upcycle---Vmlldzo4NDA3MzI2](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Noise-upcycle---Vmlldzo4NDA3MzI2)
Refer to caption: x33.png

##### Noise upcycling

For the creation of Qwen2-MoE [, , ], the authors add 50% of gaussian noise to feedforward networks before continuing training in an upcycled setup []. Komatsuzaki et al. [] also report that they experimented with adding noise but did not find it beneficial. In Figure 28, we experiment with regular upcycling versus adding noise by randomly replacing 50% of each MLP with numbers drawn from a normal distribution with a standard deviation of 0.02 following. We find that after 700 billion tokens, the no noise variant still performs slightly better but both appear to converge to the same performance. If training further, it is possible that the noise variant eventually outperforms the no noise variant, but at that point, it may make more sense to just train the MoE from scratch (§4.1.5).

Figure: Figure 29: Sharing the same MoE across layers versus a regular dense LM. The number of experts in the MoE is equivalent to its number of layers. Thus, because the MoE is shared across layers, it has the same number of total and active parameters as the dense model. More results, logs, and configurations: [https://wandb.ai/ai2-llm/olmoe/reports/Plot-Shared-vs-Dense--Vmlldzo4NDI0MTc5](https://wandb.ai/ai2-llm/olmoe/reports/Plot-Shared-vs-Dense--Vmlldzo4NDI0MTc5)
Refer to caption: x34.png

##### Shared Layer

Some work has investigated Mixture-of-Experts with weights shared across layers in the context of Universal Transformers [, , ]. We test whether layer-shared Mixture-of-Experts can beat non-shared dense models in Figure 29. The layer-shared MoE uses a load balancing loss that is applied at the model level rather than at the layer level. This gives the model more flexibility by allowing it to completely deactivate certain experts for some layers and even emulate a dense model by always activating one separate expert for each layer. This makes it a generalization of the dense model which motivated our hypothesis that it may perform better than the dense model. However, in practice, we find that both perform similarly with the regular dense models even maintaining a small advantage on validation loss and HellaSwag. One possible advantage of layer-shared MoEs is that they can allow for better load balancing at inference. If prompts come in continuously, then newly incoming prompts can be batched with previous prompts that have already passed through several layers and sent through the MoE module together, as the MoE module is the same regardless of whether it is the first or last layer. Sharing also reduces throughput by around 20% during training, which further motivates our decision not to use it for OLMoE-1B-7B.

##### KTO experiments

In Table 14 we experiment with the number of steps (5,000 vs. 10,000) and the optimizer (Adam [] vs. RMS) used for KTO []. Based on these experiments we use the RMS optimizer and the checkpoint at 5,000 steps in §4.3.

**Table 14: KTO adaptation experiments. 5,000 and 10,000 steps correspond to 1.3 and 2.6 epochs on our adaptation dataset (§2), respectively.**
|  |  |  |  | Human- | Alpaca- |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Task () | MMLU | GSM8k | BBH | Eval | Eval 1.0 | XSTest | IFEval | Avg |
| Setup () | 0-shot | 8-shot CoT | 0-shot | 0-shot | 0-shot | 0-shot | 0-shot | 0-shot |
| Metric () | EM | EM | EM | Pass@10 | %win | F1 | Loose Acc |  |
| KTO, 5,000 steps, RMS | 51.2 | 45.5 | 34.1 | 57.1 | 81.6 | 86.6 | 47.5 | 57.7 |
| KTO, 10,000 steps, RMS | 51.0 | 41.0 | 34.7 | 53.8 | 81.0 | 62.3 | 47.5 | 54.2 |
| KTO, 5,000 steps, Adam | 51.2 | 42.0 | 35.3 | 55.6 | 81.0 | 84.5 | 46.6 | 56.0 |
| KTO, 10,000 steps, Adam | 51.0 | 43.0 | 34.1 | 54.9 | 79.7 | 62.7 | 47.5 | 53.3 |

## Appendix G Additional Analysis

Figure: Figure 30: Vocabulary specialization for OLMoE-1B-7B when considering all 8 activated experts. Equivalent to $k=8$ in Equation 8.
Refer to caption: x35.png

Figure: Figure 31: Vocabulary specialization for Mixtral-8x7B when considering all 2 activated experts. Equivalent to $k=2$ in Equation 8.
Refer to caption: x36.png

Figure: Figure 32: Vocabulary specialization across domains of OLMoE-1B-7B (top) and Mixtral-8x7B (bottom). We visualize how often token IDs get routed to specific experts. We only include IDs that appear at least 8 times in the various corpora. Vertical gray lines correspond to uniform routing (8/64=12.5% for OLMoE-1B-7B as it has 64 experts, 8 of which are activated; 2/8=25% for Mixtral as it has 8 experts, 2 of which are activated). For example, among all token IDs in GitHub that get routed to Expert 0 at least 8 times for OLMoE-1B-7B, 40% of them get routed to Expert 0 with a probability of 100% (upper left) indicating that Expert 0 is specialized on those token IDs. For OLMoE-1B-7B there is much frequency at the routing probability extremes (0% or 100%) indicating that these experts exclusively focus on certain token IDs, especially for specific domains (§5.3) like GitHub and arXiv.
Refer to caption: x37.png

Figure: Figure 33: Load imbalances in selective layers after adaptation. We visualize how often tokens from our instruction tuning dataset (§2) get routed to the 8 active experts out of the 64 total experts ($k=1$ in Equation 7). Horizontal gray lines correspond to uniform routing (8/64=12.5% per expert). Although we run SFT and DPO without loss balancing loss (§4.3), we observe that the load distribution does not change substantially.
Refer to caption: x39.png

Figure: Figure 34: Domain specialization of OLMoE-1B-7B (top) vs. Mixtral-8x7B (bottom) of the top-$1$ routed expert. We visualize how often tokens from different domains get routed to the 64 (OLMoE) or 8 (Mixtral) experts at the end of pretraining. Unlike in Figure 22, here we only consider tokens routed to the top-1 expert ($k=1$ in Equation 7). Horizontal gray lines correspond to uniform routing (1/64=1.56% per expert for OLMoE-1B-7B and 1/8=12.5% for Mixtral).
Refer to caption: x40.png

Figure: Figure 35: OLMoE-1B-7B token routing across layers. We visualize how often tokens from different domains get routed to a pair of experts across layers under top-1 routing, corresponding to Figure 34. The size of each rectangle is proportional to the total number of tokens an expert receives, while the flow between two experts shows the proportion of tokens routed to both experts. We only show experts that receive tokens 50% above random chance and use stronger coloring for larger flows. We observe some instances of cross-layer coordination between pairs of experts, e.g., expert 27 in layer 7 and expert 57 in layer 15 process a substantial fraction of Wikipedia tokens together. The flows between layers 0 7 and 7 15 are independent in this visualization.
Refer to caption: x42.png

Figure: Figure 36: Mixtral-8x7B token routing across layers. We visualize how often tokens from different domains get routed to a pair of experts across layers under top-1 routing, corresponding to Figure 34. The size of each rectangle is proportional to the total number of tokens an expert receives, while the flow between two experts shows the proportion of tokens routed to both experts. The flows between layers 0 7 and 7 15 are independent in this visualization.
Refer to caption: x46.png

## Appendix H Limitations and Future Work

We highlight four key limitations with this release of OLMoE-1B-7B. We look forward to addressing these issues in future iterations of OLMoE.

##### More parameters

OLMoE-1B-7B has 7B total parameters out of which 1B are activated for each input token. This small size makes OLMoE-1B-7B very cheap to use, yet we demonstrate in this work that it outperforms much more expensive models (Figure 1). However, using only 1B parameters for each input token also limits the capabilities of OLMoE-1B-7B as seen by its performance compared to models that use $>$7 more parameters, such as Llama3.1-8B in §3. While it may be possible that more parameters are not needed to match 8B models and beyond [], in the short-term adding parameters is an easy way to improve the performance of OLMoE, at least allowing the model to utilize more than 1B parameters per input, possibly via recursion [] or agentic workflows [, ]. Relatedly, changing the allocation of parameters to e.g. vocabulary versus non-vocabulary parameters is another avenue for improvement [].

##### More data

We train OLMoE-1B-7B for 5 trillion tokens, however, some recent dense models train significantly longer, such as Llama 3 with 15 trillion tokens []. To the best of our knowledge, there has been no large MoE that has been overtrained [] as much as OLMoE-1B-7B. Specifically, taking the active parameters of OLMoE-1B-7B, our token multiplier [] is around 5,000 (5T / 1B). There are likely benefits to training even longer, but to what degree overtraining is effective for MoEs and how it differs from dense models still requires more research [].

##### Multimodal

OLMoE-1B-7B is a text-only large language model, thus it cannot take inputs or produce outputs in other modalities like images or audio. This limits its utility for the large variety of multimodal use cases of such models [, , , , , , , , ]. There has been early work on open multimodal MoEs [, , , , , ] and we look forward to making future versions of OLMoE a part of that.

##### Multilingual

We pretrain OLMoE-1B-7B on a predominantly English corpus and exclusively evaluate on English tasks. This may severely limit the usefulness of our model for research on non-English language models [, , , , , ]. While there has been work on training language-specific LMs [, ], it is more likely that as we add more data to build better future iterations of OLMoE we will mix in more non-English data due to data constraints []. This may make future OLMoE models perform better in non-English languages.

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: See Above.
soundness: 4
strengths: - The writing in this paper is clear and easy to follow.
- The paper advances MoE research by providing a fully open-sourced, state-of-the-art MoE architecture, which is beneficial for the research community.
- The paper presents a thorough analysis of key design choices in MoE, offering valuable guidance on building high-performance MoE models.
- The analysis is insightful, with discussions on phenomena such as router saturation and expert co-activation providing fresh perspectives and meaningful implications for the field.
summary: This paper introduces OLMoE, a fully open, state-of-the-art language model built on a sparse Mixture-of-Experts (MoE) architecture. The authors conducted extensive experiments to validate the effectiveness of the proposed method, including evaluations after pre-training and adaptation phases. Additionally, they explored key design choices within the MoE framework, examining factors like expert granularity, routing strategies. Their analyses provided valuable insights into MoE, including router saturation, expert co-activation, and domain/vocabulary specialization.
weaknesses: I have a question regarding the experimental results: were the model parameters quoted directly from the original paper for the results shown in Table 2? For instance, in the original paper, OpenMOE’s activation parameter count is reported as 2.1B, whereas Table 2 shows an activation parameter count of 2.9B for OpenMOE. I recommend that the authors carefully verify the accuracy of these values.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
questions: The work is well presented and possible suggestions for improvements are addressed in the future work section.
soundness: 4
strengths: 1) Strong empirical results with state-of-the-art performance for 1B active parameters.
2) Good exploration of the MoE design space which forms a good guide for MoE model design.
3) Novel analysis of routing behavior in MoE models during training and inference.
4) This is the only MoE model where the model weights, code, data and checkpoints are openly available and thus the work is entirely reproducible.
summary: This paper presents a mixture-of-experts (MoE) LLM model called OLMoE that has 1B active parameters and 7B total parameters. The OLMoE model Pareto dominates many state-of-the-art models in the performance vs. active parameters space. The paper explores and presents insights on what is optimal in the design-space of MoE parameters and present analysis of routing behavior in MoEs.
weaknesses: 1) Other state-of-the art MoE models in related works are not exactly in the same parameter count configuration (1B/7B) so an exact comparison cannot be made to this model's performance.
2) Most of the design choices and training choices are based on prior work and the novelty is more in the design space exploration and analysis of routing behavior.

---

confidence: 5
contribution: 4
initial_rating: 10
presentation: 4
questions: 1) What do you think about the necessity of expert parallelism? This model used dropless MoE, so it anyway will be unbalanced when using expert parallelism during training and inference. Without expert parallelism, it is still okay when the model is small. However, if we are aiming at a very large model, which has very large experts even if we are using the "fine-grained MoE", the expert parallelism would still be required? So how can we handle the token drop problem in this case?
soundness: 4
strengths: 1) There is no doubt that training MoE LLMs is challenging. This work offers a couple of important takeaways about how to train a good MoE LLMs, which is very helpful to the community.
2) The presentation is very clear. For instance, the Table 1 delivers many key designs clearly at the early section of the paper.
3) The model performance is good as well. As shown in Table 2 and 3, the model performs competitive with dense open models and partially open models (e.g. Qwen, Deepseek).
4) The Analysis in Section 5 is informative, which greatly help readers and authors to understand how is the model working. This can also greatly speedup the growth of the community.
summary: This work is devoted to sharing the insights, data, and checkpoints of a series of MoE LLMs. The model achieved promising results on various benchmarks as a fully open model family.
weaknesses: 1) Although the model has been relatively large, it is still much smaller than the SoTA MoE LLMs. I understand it is hard to get enough training resource for a fully open projects.

Meta-review: The paper introduces a state-of-the-art language model leveraging a sparse Mixture-of-Experts (MoE) architecture. It presents novel findings on MoE training, defines and analyzes new routing properties showing high specialization in their model, and open-sources all aspects of their work, including model weights, training data, code, and logs.

The paper provides significant contributions to the field of MoE models, offers valuable insights, and is fully open-source, making it a valuable resource for the research community. Based on the reviews and the authors' thorough responses, I recommend acceptance.

Decision: Accept (Oral)