2 OLMo 2 Furious β€” part 3/7

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

Since fast spikes are difficult to understand with contemporary graphing tools, we compute a spike score as an objective measure. Concretely, we define the spike score as the percentage of values in a time series that are at least seven standard deviations away from a rolling average of the last 1,000 1,000 values 8 8 8 Spike score is conceptually similar to spike mitigation proposed by Karpathy (2024) . . We use spike score primarily on training loss and L2 norm of the gradient, but the measure can be computed on any time series.

##### Empirical results

To experiment with model initialization, we first create a baseline run that reproduces spikes quickly. We do so by mainly reducing the warmup period. The effect was immediate and dramatic (Figure 4 ), and persists across model scales and token counts. In our ablation, the new initialization had no loss spikes, and the spike score for the L2 norm of the gradient went from 0.40 0.40 to 0.03 0.03 . The new initialization converges slightly slower; we make up for this difference by improving other hyperparameter settings (Section Β§ 3.4 ).

### 3.3 Architecture Improvements

#### 3.3.1 Nonparametric layer norm and RMSNorm

OLMo 2 uses RMSNorm, which is standard in most transformer implementations. OLMo-0424 used a nonparametric layer norm for performance and to work around bugs in the libraries we were using, but by the time we developed OLMo 2 , the bugs were no longer an issue, the hardware was faster, and we wanted to settle on a safe approach. Our ablations show no difference between the two, so we switch back to RMSNorm.

#### 3.3.2 Reordered norm and QK-norm

Figure 7 shows the effect of applying the layer normalization to the outputs of the MLP and attention blocks instead of the inputs. We further apply another normalization, also RMSNorm, to the queries and keys in the attention block. In isolation, neither of these changes yield good results, but together they improve both the growth and the spikiness of the L2 norm of the gradient. The following table summarizes the difference in the location of the layer normalization:

OLMo-0424 OLMo 2 𝒉 := 𝒙 + Attention ​ ( LN ​ ( 𝒙 ) ) \boldsymbol{h}:=\boldsymbol{x}+\text{Attention}(\text{LN}(\boldsymbol{x})) 𝒉 := 𝒙 + RMSNorm ​ ( Attention ​ ( 𝒙 ) ) \boldsymbol{h}:=\boldsymbol{x}+\text{RMSNorm}(\text{Attention}(\boldsymbol{x})) 𝒉 out := 𝒉 + MLP ​ ( LN ​ ( 𝒉 ) ) \boldsymbol{h}_{\text{out}}:=\boldsymbol{h}+\text{MLP}(\text{LN}(\boldsymbol{h})) 𝒉 out := 𝒉 + RMSNorm ​ ( MLP ​ ( 𝒉 ) ) \boldsymbol{h}_{\text{out}}:=\boldsymbol{h}+\text{RMSNorm}(\text{MLP}(\boldsymbol{h})) 𝒙 \boldsymbol{x} is the input to the layer, 𝒉 \boldsymbol{h} is an intermediate hidden state, and 𝒉 out \boldsymbol{h}_{\text{out}} is the output.

Liu et al. (2021) first introduced layer norm the idea of reordering layer norm. It was subsequently picked up by Chameleon Team (2024) . QK-norm was first developed in Dehghani et al. (2023a) .

#### 3.3.3 Z-Loss

Following Chowdhery et al. (2022) , Chameleon Team (2024) , and Wortsman et al. (2023) , we apply z-loss regularization by adding 10 βˆ’ 4 β‹… log 2 ⁑ Z 10^{-4}\cdot\log^{2}Z to our loss function, where Z Z is the denominator in the softmax over the logits. This discourages the activations in the final softmax from growing too large, improving the stability of the model.

Figure 8 shows a stark difference between the z-loss implementation of the popular Flash Attention library ( Dao, 2024 ) , and an implementation using only Python primitives. Apart from the attention mechanism it is known for, Flash Attention also provides an optimized implementation of cross-entropy loss, which includes a version of z-loss. To retain flexibility in settings that are not compatible with Flash Attention, we have a separate implementation written in PyTorch. Both implementations produce the same result in the forward pass, but exhibit different behavior in the backward pass. We suspect the root cause lies in differences in precision. In our experiments, this does not affect cross entropy loss during training, or the model’s performance on downstream tasks. However, out of an abundance of caution we abandon the fork with custom z-loss implementation and re-train from the original point of divergence. During a training run we cannot switch implementations safely, so we avoid doing so as much as possible.

### 3.4 Hyperparameter Improvements

#### 3.4.1 Ο΅ \mathbf{\epsilon} in AdamW

Figure 9 shows the result of decreasing the AdamW Ο΅ \epsilon from 10 βˆ’ 5 10^{-5} to 10 βˆ’ 8 10^{-8} . 10 βˆ’ 8 10^{-8} is the default in PyTorch, but some popular LM training code bases come with a default of 10 βˆ’ 5 10^{-5} . The lower value allows for larger updates early in training, and helps the model learn faster during a period where we’ve typically seen a lot of instability. As a result, the gradient norm settles much more quickly and remains permanently lower.

#### 3.4.2 Weight decay on embeddings

Figure 10 shows the change in training dynamics following a decision to exclude weight decay for embeddings. OLMo uses a standard formulation of weight decay, where every parameter is multiplied by 1 βˆ’ ( 0.1 β‹… l ​ r ) 1-(0.1\cdot lr) at every step. This regularization term discourages parameters from growing too large, but in the case of token embeddings it overshoots the mark and results in very small embeddings. As discussed by Takase et al. (2024) , small embeddings can produce large gradients in early layers because the Jacobian of 𝗅𝖺𝗒𝖾𝗋 ​ _ ​ π—‡π—ˆπ—‹π—† ​ ( x ) \mathsf{layer\_norm}(x) w.r.t. x x is inversely proportional to β€– x β€– \left\lVert x\right\rVert , and, in early layers, the norm of the residual stream is essentially the norm of the embeddings. We experiment with the full range of remedies discussed in Takase et al. (2024) , but found that they impacted the speed of convergence. Instead, we simply turn off weight decay for embeddings and observe that embedding norms settle in a healthy region as training progresses.

## 4 Deep Dive: Mid-training Recipe

Recent works have suggested that a multi-stage approach to base model training can lead to measurable improvements in capabilities ( Blakeney et al., 2024 ; Ibrahim et al., 2024 ; Feng et al., 2024 ) . In previous OLMo iterations, we also found that both learning rate schedule ( OLMo 1 ; Groeneveld et al. 2024 ) and data mixture ( OLMo-0424 ; Ai2 2024 ) play an important role. We refer to interventions at this stage of model development as mid-training 9 9 9 while the concept of chaining of multiple stages of self-supervised training is not new ( e.g. , Gururangan et al. 2020 ), we trace the use of mid-training to Abdin et al. (2024a) and OpenAI (2024) . .

From afar, our approach is simple: after the pretraining stage, we generate domain-specific data mixtures and restart training, linearly driving the learning rate down to zero. Our goal is to imbue specialized knowledge and improve capabilities; feedback on these improvements comes from key benchmarks, such as math-specific tasks such as GSM8K.

### 4.1 Learning rate annealing

Our starting point for learning rate experiments was the setting from Grattafiori et al. (2024) . To initialize the optimizer state for the 7B variant, we linearly warm up the learning rate to its peak of 3 β‹… 10 βˆ’ 4 3\cdot 10^{-4} over the first 2000 steps. Then, we use a standard cosine decay over 5T tokens. Previous experience with OLMo-0424 suggests that the last part of a cosine decay schedule can be cut off and replaced by a linear decay to zero with little loss of performance. Accordingly, for the 7B variant, we stop the schedule at 4T tokens and then switch to mid-training as described in Section Β§ 4 . The 13B ran with a higher peak learning rate from the start, so we decided to run it to 5T tokens before moving to the mid-training stage.

Figure 11 shows different runs with four additional learning rate values: 6 β‹… 10 βˆ’ 4 6\cdot 10^{-4} , 9 β‹… 10 βˆ’ 4 9\cdot 10^{-4} , 12 β‹… 10 βˆ’ 4 12\cdot 10^{-4} , and 30 β‹… 10 βˆ’ 4 30\cdot 10^{-4} . In particular, we tried double, triple, quadruple, 10 Γ— \times , and 30 Γ— \times the original learning rate. The last, 30 β‹… 10 βˆ’ 4 30\cdot 10^{-4} , showed training instabilities already during learning rate warm-up, with several loss spikes that did not recover fully, so we abandoned this variant quickly. The other values trained normally and showed an interesting pattern. Looking purely at training loss, higher learning rates universally perform better early on (as long as they avoid loss spikes), but eventually the lower learning rate setting overtakes the others (Figure 11 ). Notably, when comparing 3 β‹… 10 βˆ’ 4 3\cdot 10^{-4} and 6 β‹… 10 βˆ’ 4 6\cdot 10^{-4} , the cross-over point is well past 200B tokens. A shorter hyperparameter experiment might come to the wrong conclusion.

One of the motivations for this line of experimentation was to find out whether a higher learning rate would make the annealing step more effective. The conjecture is that the worse training loss during pretraining is compensated for when the learning rate decays to zero. To test this hypothesis, we took a checkpoint from each of our four variants after 300B tokens, and decayed the learning rate to zero over 50B tokens. To account for the possibility that the effect of higher learning rates needs more steps to unfold, we tried the three higher settings and decayed the learning rate over 100B tokens, for a total of seven experiments. The results show that a higher learning rate does make mid-training more effective, but it does so by exactly the amount that the pretraining is worse. All four variants show the same training loss at the end of the procedure, though the lowest setting lags behind the others by a small amount.

Table 8 shows that the result is consistent for longer training runs as well. We took two variants, 3 β‹… 10 βˆ’ 4 3\cdot 10^{-4} and 6 β‹… 10 βˆ’ 4 6\cdot 10^{-4} , and repeated the experiment after training for 1T and for 2T tokens. We chose these variants because 3 β‹… 10 βˆ’ 4 3\cdot 10^{-4} is the baseline from Grattafiori et al. (2024) , and 6 β‹… 10 βˆ’ 4 6\cdot 10^{-4} showed, by a slim margin, the best training loss. Our results show virtually no difference between the two settings, both on training loss and a mix of nine downstream tasks from the OLMES suite ( Gu et al., 2024 ) shown in Table 8 . Evaluating the models on downstream tasks is noisier, but mirrors the findings based on training loss only.

Finally, we wanted to see if a higher learning rate during the pretraining stage would result in a more effective mid-training stage when switching to higher quality data. To match our training setup as much as possible within the available compute budget, we took the same two settings ( 3 β‹… 10 βˆ’ 4 3\cdot 10^{-4} and 6 β‹… 10 βˆ’ 4 6\cdot 10^{-4} ), and linearly decayed the learning rate to 0 over 100B high quality tokens. Once again, the results show little difference. The final scores on the OLMES evaluation suite are within 0.1 points of each other. However, looking at other metrics may still reveal a meaningful difference between the two settings. The mix of high quality tokens targets math specifically, and on GSM8K (which is not part of the OLMES suite), the high learning rate setting is 2.8 points better than the lower learning rate. More study is needed to turn this interesting data point into a dependable result.

This finding contradicts machine learning folk wisdoms such as β€œhigher learning rates are always better” or β€œarea under the learning curve matters” ( McCandlish et al., 2018 ) . It expands on Wortsman et al. (2023) , who observed that smaller models’ performance is largely invariant to learning rate over several orders of magnitude when trained to the end of a cosine schedule, and further found that QK-norm (section 3.3.2 ) and z-loss (section 3.3.3 ), which we use as well, enhance this effect. We find that these results still hold even at much larger scales of tokens and parameters, and, crucially for our training efforts, with our modified learning rate schedule.

Due to cost concerns we did not explore the full range of learning rates. This is the main limitation of this line of experimentation. It would be interesting to run a wider sweep of learning rates to accurately define the boundaries of the plateau we appear to be training in.

### 4.2 Data Curriculum: Dolmino Mix 1124

In this section, we describe our experimental process for curating our mid-training data. We collectively refer to the resulting dataset and mixtures created for this mid-training stage as Dolmino Mix 1124 . An overview of the contents of this dataset is provided in Section Β§ 2.4 (Table 5 ). In detail, we use the following procedure in our mid-training recipe: β€’ Identify a mix of high-quality sources to improve performance across the entire development benchmark suite (Section Β§ 4.3 ).

β€’ For patching specific capabilities (specifically, in the case of OLMo 2 , math), collect and evaluate domain-specific datasets to mix during mid-training (Section Β§ 4.4 ). We found that these sources can be independently assessed through a technique we dub microannealing (Section Β§ 4.4.2 ); their effectiveness persists when mixed with rest of sources.

β€’ Following experiments described in Section Β§ 4.1 , we mix high-quality sources and math-specific data in three different token budgets (50B, 100B, 300B). The smaller mix is used to mid-train OLMo 2 7B, while OLMo 2 13B and 32B are annealed on the larger ones. For both OLMo 2 7B, 13B and 32B, we find that averaging weights of different checkpoints trained on same mixture but different data order seeds consistently improves over individual checkpoints (Section Β§ 4.5 ). To demonstrate this on the small scale, we also include results for a 1B model that receives similar interventions as the 7B model.

Table 9 summarizes the dramatic impact of this mid-training phase on both development and held-out evals. OLMo 2 7B model improves, on average by 10.6 points, surpassing the larger 13B model after the pretraining stage. For its part, OLMo 2 13B benefits equally from mid-training, improving its average performance by 10.3 points. Both models see improvements in knowledge-intensive, multiple-choice (Arc challenge: 72.6 β†’ 79.8 72.6\rightarrow 79.8 for 7B, 80.2 β†’ 83.5 80.2\rightarrow 83.5 for 13B; MMLU: 59.8 β†’ 63.7 59.8\rightarrow 63.7 for 7B, 63.4 β†’ 67.5 63.4\rightarrow 67.5 for 13B; AGIEval: 44.6 β†’ 50.4 44.6\rightarrow 50.4 for 7B, 48.2 β†’ 54.2 48.2\rightarrow 54.2 for 13B), reading comprehension (Natural Questions: 29.0 β†’ 36.9 29.0\rightarrow 36.9 for 7B, 34.6 β†’ 46.7 34.6\rightarrow 46.7 for 13B; DROP: 40.7 β†’ 60.8 40.7\rightarrow 60.8 for 7B, 49.6 β†’ 70.7 49.6\rightarrow 70.7 for 13B), and math skills (GSM8K: 24.1 β†’ 67.5 24.1\rightarrow 67.5 for 7B, 37.3 β†’ 75.1 37.3\rightarrow 75.1 for 13B) benchmarks.

### 4.3 Dolmino Mix 1124 : High Quality Sources

Following the recipe from the previous OLMo iteration ( Ai2, 2024 ) , we start by curating a higher quality subset of pretraining mix, and expand it with more academic and encyclopedic material. In particular, we consider the following sources (summarized in Table 10 ):

##### High quality web

To filter the web subset used in pretraining, we experiment with two existing quality classifiers: β€’ FastText classifier from Li et al. (2024) . To train this model 10 10 10 mlfoundations/fasttext-oh-eli5 , Li et al. sampled positive documents from the Reddit subset in ELI5 ( Fan et al., 2019 ) , and demonstrations from Open Hermes 2.5 11 11 11 datasets/teknium/OpenHermes-2.5 . Negatives are sampled at random from the DCLM pipeline.