2 OLMo 2 Furious — part 2/7

Cleaned-text code points [15735, 31320). Bibliography omitted here; source.md includes it.

We curated a smaller, focused mixture— Dolmino Mix 1124 —to imbue the model with domain knowledge from increased exposure to STEM references and high quality text as well as skills that remained lacking after the initial pretraining stage (e.g. math-solving capabilities). We up-sample high-quality web documents and curated non-web sources; we also employ synthetic data crafted to patch math capabilities of the model.

##### Model Merging or “Souping”

To get the most out of this high-quality data, and to find a better local minimum, we perform this step multiple times with different random data orders, and then average the resulting models ( Matena and Raffel, 2022 ; Wortsman et al., 2022 ) . For OLMo 2 7B, we anneal three separate times for 50B tokens each, with different randomized data orders; we average the resulting models to produce the final model. For both OLMo 2 13B and OLMo 2 32B, we train three separate times for 100B tokens each (same number of update steps as the 7B), and then a fourth time for 300B tokens. The final model is the average of all four models. For further details, refer to Section § 4 .

##### Overall

In total, OLMo 2 7B is trained on 4.05 4.05 trillion tokens ( 3.90 3.90 trillion for pretraining stage), OLMo 2 13B is trained on 5.6 5.6 trillion tokens ( 5 5 trillion for pretraining stage), and OLMo 2 32B is trained on 6.6 6.6 trillion tokens ( 6.06 6.06 trillion for pretraining stage).

### 2.4 Base Model Data

We provide a brief overview of the data mix for pretraining and mid-training in this section.

#### 2.4.1 Pretraining data: OLMo 2 Mix 1124

The mix used for this stage is shown in Table 4 . It consists of approximately 3.9 trillion tokens, with over 95% derived from web data. We refer to this set as OLMo 2 Mix 1124 . This is the same pretraining data used in OLMoE ( Muennighoff et al., 2024 ) : We combine data from DCLM ( Li et al., 2024 ) and Dolma 1.7 ( Soldaini et al., 2024 ) . From DCLM, we use the “ baseline 1.0 ’’ mix. 5 5 5 Available at mlfoundations/dclm-baseline-1.0 From Dolma, we use the arXiv ( Together AI, 2023 ) , OpenWebMath ( Paster et al., 2023 ) , Algebraic Stack, peS2o ( Soldaini and Lo, 2023 ) , and Wikipedia subsets. arXiv, OpenWebMath, and Algebraic Stack were originally part of ProofPile II ( Azerbayev et al., 2023 ) . Finally, we include code from StarCoder ( Li et al., 2023b ) , which is derived from permissively-licensed repositories from GitHub ( Kocetkov et al., 2022 ) . In an attempt to include higher quality code, we remove any document from a repository with fewer than 2 stars on GitHub. Further, through manual inspection of this source, we found it to contain documents encoded in binary format or containing mostly numerical content; to remove them, we discarded documents whose most frequent word constitutes over 30% of the document, or whose top-2 most frequent words constitute over 50% of the document. To mitigate possible training loss spikes, we remove documents with repeated sequences of 32 or more n-grams. We report details and show effectiveness of this intervention in Section § 3.1 .

#### 2.4.2 Mid-training data: Dolmino Mix 1124

After the initial pretraining stage on mostly web data, we further train with a mixture of web data that has been more restrictively filtered for quality and a collection of domain-specific high quality data, much of which is synthetic. The purpose of this mixture is to imbue the model with math-centric skills and provide focused exposure to STEM references and high quality text. We generate several variants of this mixture, with varying sizes, but generally refer to this mixture as Dolmino Mix 1124 . The base sources from which Dolmino Mix 1124 is subsampled are described in Table 5 . We refer the reader to Section § 4 for a deep dive detailing our processes for experimenting and curating data for this mix.

### 2.5 Evaluation and Results

OLMo 2 is evaluated via standard language model benchmarks. Further, we apply post-training to OLMo 2 and evaluate the result— OLMo 2-Instruct —on a diverse set of tasks to assess the adaptation potential of our base model.

##### Base Model Evaluation:

We evaluated OLMo 2 and other baseline models using the OLMES evaluation suite ( Gu et al., 2024 ) , which includes a range of benchmark datasets for both multiple-choice and generative tasks, using standardized prompts and in-context examples for few shot predictions. Full descriptions of benchmark tasks in Appendix A.1 . For multiple-choice tasks, we evaluate accuracy; for generative tasks, we evaluate F1 to account for partial matches. Additionally, to avoid overfitting our recipe to these benchmarks, we maintained a held-out suite of tasks which were not used for model development decisions; we advocate for a standard practice of declaring development vs held-out evaluation tasks for model developers. 6 6 6 GSM8k ( Cobbe et al., 2021 ) was only partially held-out, as we subsampled 200 of 1319 GSM8k examples for mid-training data development when we noticed poor math capabilities after pretraining; we call this dev set GSM ∗ . The remaining 1119 GSM8k examples we reserve as held-out and report final performance on them only.

Table 6 contains overall results. We find our OLMo 2 models are competitive with the best open-weights models of comparable size, despite OLMo 2 requiring far fewer training FLOPs (see Figure 1 ) and maintaining full openness (e.g. training data) . We find that gains observed on development metrics largely translate to our unseen evaluation suite, indicative of a generalizable training recipe.

Overall, we find that gains observed on development metrics largely translate to our unseen evaluation suite. Of course, we have no guarantee that tasks we consider unseen during development of OLMo 2 are not part of the development set of other models we compare. Nevertheless, we think it should be standard practice for model developers to keep a subset of evaluation tasks unseen and to declare which these are, in technical reports. Further, we encourage other open-weight model developers to clearly state which tasks are being monitored during model development.

##### Post-Training Recipe and Evaluation

For post-training we apply our Tülu 3 ( Lambert et al., 2024 ) recipe with supervised finetuning, on-policy preference tuning, and reinforcement learning with verifiable rewards (RLVR). 7 7 7 We made minor modifications to the preference data to use generations from permissively-licensed models and added a multi-stage RLVR training protocol to optimize final performance, but otherwise followed the recipe as-is. The resulting models— OLMo 2-Instruct —are evaluated in Table 7 on general and precise instruction following, math, knowledge reasoning, and safety tasks from the same evaluation suite used by Lambert et al. (2024) . Full descriptions of benchmark tasks in Appendix A.2 .

Table 7 contains downstream results. We find OLMo 2-Instruct models are competitive with the best instruction-tuned open-weights models and even some popular proprietary models . This shows the usefulness of OLMo 2 as a powerful base model that serves as an excellent starting point for fully open post-training research. Full post training details are in Section § 5 .

## 3 Deep Dive: Pretraining Stability

While OLMo-0424 achieved performance within expected ranges for its compute budget, the training dynamics were characterized by a couple of concerns: • Sudden spikes in the loss, and more frequently, in the gradient norm during training. In experiments, we found that increasing model size increased the frequency of spikes. Furthermore, our experiments revealed that more dramatic spikes in gradient norm often preceded training loss spikes.

• Slow growth in the magnitude of the gradient norm over the training run. This was correlated with increasing frequency of spikes in the gradient norm (and training loss).

Ultimately, a combination of these issues would lead to training divergence, making training at larger scales impossible. This situation motivated our training stability investigation into the causes of these issues and their mitigations. Figure 2 shows our training curves before and after implementing our mitigations, which we summarize below: • Repeated n-grams: We filter pretraining data to remove repeated n-grams in pretraining data, as they can lead to loss spikes (§ 3.1 ).

• Initialization: We switch from scaled initialization ( Zhang et al., 2019 ) to initializing all parameters with a mean of 0 and a standard deviation of 0.02 (§ 3.2 ).

• RMSNorm: We use the RMSNorm variant of LayerNorm to normalize activations instead of non-parametric LayerNorm (§ 3.3.2 ).

• Reordered norm: We normalize the outputs to the attention and feed-forward (MLP) layers within each transformer block instead of the inputs (§ 3.3.2 ).

• QK-norm: We normalize the key and query projections with RMSNorm before calculating attention (§ 3.3.2 ).

• Z-Loss: We adopt z-loss regularization, a regularization term that keeps final output logits from growing too large (§ 3.3.3 ).

• Weight decay: We exclude embeddings from weight decay (§ 3.4.2 ).

• ϵ \epsilon in AdamW: We lower the ϵ \epsilon of AdamW from 10 − 5 10^{-5} to 10 − 8 10^{-8} (§ 3.4.1 ).

In the following, we will discuss the experiments and results that led us to these interventions. We compare our revised strategies with OLMo-0424 , the most recent version of OLMo with fully-open model weights, data, and documentation.

### 3.1 Repeated n-Grams

Data can be a cause of both gradient norm and loss spikes. When investigating training batches at which spikes occurred, we found a high prevalence of instances containing long, repeated n-gram sequences. Here are three examples of such sequences:

g4ODg4ODg4ODg4ODg4ODg4ODg4ODg4ODg4ODg4ODg4ODg4ODg4ODg4OD... [\n 365, 0, 667, 1000, 1000, 667, 667, 667, 667, 667, ... ’ 255, 255, 255, 255, 255, 255, 255, 255, \n255, 255, ...

In a series of experiments, we found these sequences are often associated with spikes, though we note that this relationship is not deterministic: • The same n-gram sequence may spike for a larger model but not for a smaller model trained on the same data.

• The same n-gram sequence may spike for one data training ordering, but not after the data is reshuffled.

• The same n-gram sequence associated with a spike can also be found elsewhere in training batches that did not spike.

Nevertheless, we have found evidence that broad removal of such sequences across training decreases the frequency of spikes, on average. At data curation time (Section § 2.4 ), 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. We also implement an additional safeguard in the trainer that detects these sequences during data loading and masks them when computing the loss. Figure 3 shows the effect of masking the loss of input sequences containing repeated n-grams. This intervention results in a clear mitigation—though not complete elimination—of gradient spikes. It had no effect on the slow growth in gradient norm.

### 3.2 Model Initialization

Figure 4 shows the improvement to training stability from OLMo 2 ’s initialization scheme. In OLMo 2 , we initialize every parameter from a normal distribution with a mean of 0 and a standard deviation of 0.02. In contrast, OLMo-0424 ’s initialization, first suggested in Zhang et al. (2019) and implemented by Gururangan et al. (2023) , scaled input projections by 1 / d model 1/\sqrt{d_{\mathrm{model}}} , and output projections by 1 / 2 ⋅ d model ⋅ layer ​ _ ​ idx 1/\sqrt{2\cdot d_{\mathrm{model}}\cdot\mathrm{layer\_idx}} at every layer. In other words, later layers were initialized to smaller values.

We perform several analyses to study the impact of initialization, showing that OLMo 2 ’s initialization is superior to OLMo-0424 initialization. Our empirical analysis suggests it better preserves the scale of activations and gradients across layers, allowing deep models to be trained more stably, and it exhibits properties associated with hyperparameter transfer across models of different widths. These two properties together give us confidence that deep models will train stably and that the initialization hyperparameters of our smaller models could transfer to larger scales.

##### Gradient and activation growth

A fundamental concern for training deep networks is ensuring that the activations and gradients do not blow up or vanish across layers, causing learning to become unstable or stagnate. Rather, we want the scale of the activations and gradients to remain roughly the same from layer to layer. Inspired by recent related work ( Cowsik et al., 2024 ) , we evaluate different candidate initializations in terms of how they affect the 2-norm of the activations and gradients across layers. Concretely, we randomly initialize a model, pass 50 random documents from The Pile ( Gao et al., 2021 ) through it, and collect the activations and gradients (of loss with respect to the activations) at the initial and final layers (ignoring embeddings). We then average these tensors across documents and time steps to get vectors 𝒗 \boldsymbol{v} at the initial layer and 𝒗 ′ \boldsymbol{v^{\prime}} at the final layer, both of length d model d_{\mathrm{model}} . Finally, we compute the following measure of expansion or contraction across layers, which we call the growth exponent : λ = 1 n layers ​ log ⁡ ( ‖ 𝒗 ′ ‖ ‖ 𝒗 ‖ ) \lambda=\frac{1}{n_{\mathrm{layers}}}\log\left(\frac{\left\lVert\boldsymbol{v^{\prime}}\right\rVert}{\left\lVert\boldsymbol{v}\right\rVert}\right) We compute λ \lambda for both the activations and gradients. Ideally, both λ \lambda ’s remain near 0, indicating that the activations and gradients do not explode or vanish across layers. Figure 5 plots the growth exponents for different randomly initialized models as a function of their widths (4096 corresponds to a full 7B model). Crucially, the growth exponent for OLMo 2 is closer to 0 than for OLMo-0424 across model widths. This suggests the OLMo 2 initialization will be more stable when training deep models in low precision, as both the activations and the gradients are more resistant to exploding or vanishing across layers compared to the original OLMo-0424 initialization.

##### Hyperparameter transfer across width

Another appealing property of the new initialization is that it scales the activation and gradient norms with width ( d model d_{\mathrm{model}} ) in a way that has been argued theoretically to be important for hyperparameter transfer across different widths. Specifically, Yang et al. (2024b) suggest that a sufficient condition for hyperparameter transfer across width is that the magnitude of each activation scalar value and its update (learning rate times gradient) remain fixed as width increases. Equivalently, the norms of the activations and their update vectors should positively correlate with d model \sqrt{d_{\mathrm{model}}} . We plot the activation and gradient norms at initialization against d model \sqrt{d_{\mathrm{model}}} in Figure 6 . Crucially, the gradient norm is more positively correlated with d model \sqrt{d_{\mathrm{model}}} for OLMo 2 compared to OLMo-0424 . Combined with Yang et al. (2024b) , this suggests that, with an initial learning rate independent of model width, the new OLMo 2 initialization will transfer better across different model widths compared to the OLMo-0424 initialization.

##### Spike score