2 OLMo 2 Furious — part 6/7

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

Beaker workloads are containerized, providing some isolation guarantees. This allows OLMo 2 workloads to run simultaneously with other jobs on the same cluster, each with unique environments and dependencies, with minimal conflicts. Notably, the Beaker executor allocates host resources in a fashion that minimizes (but doesn’t completely avoid) performance problems caused by noisy-neighbors. Containers further capture software dependencies and the runtime details of workloads. This helps run repeatable experiments, and makes it possible to replay old results even months after they happened. This stands in contrast to the more common Slurm-based setup where all workloads, whether they relate to OLMo 2 or not, share the same underlying operating system, CUDA libraries, and environment resulting in instability that makes experiments unreproducible after system changes.

Beaker also made it possible for us to take advantage of new compute sources that became available throughout the evolution of the project. Its operational simplicity made it possible for a small team of operators to quickly onboard new sources of compute.

### 6.3 Stability and Operations

Both clusters required an initial testing and burn-in period, during which we discovered and remedied problems ranging from ill-seated cables to an improper ordering of the compute nodes in the NCCL library. These periods required close collaboration with the respective hardware vendors, and both were indispensable during this process. After this period, both clusters operate approximately at the same level of reliability.

##### GPU health checks

Beaker executes a simple program prior to running workloads on the assigned GPUs. The program attempts to multiply two tensors. When failures occur, Beaker cordons the associated host and reschedules the workload, quarantining the errant node before introducing instability. This helped reduce interruptions requiring manual attention, and made it viable to configure training jobs to simply restart themselves when encountering an error, safe in the knowledge that they would be moved to the working set of compute nodes.

##### Cordoning

Beaker supports cordoning nodes as an override for the automatic health checks. A cordoned node is removed from scheduling and gets flagged for repair. In this way, Beaker effectively crowdsources the identification of bad nodes among all the users of the cluster.

##### Active monitoring

Beyond these two methods, Beaker performs industry-standard monitoring and automatic alerting based on cluster telemetry. The team has operational processes in place for responding to system issues, enabling it to resolve issues promptly.

### 6.4 Maximizing hardware utilization

Ai2’s hardware infrastructure (§ 6.1 ) has to be complemented by good model training software that gets the most out of the available resources. Increased efficiency not only lets us train larger models for more tokens, but it also improves the environmental impact of model training (§ 6.5 ), and raises experimental velocity. Further, OLMo is not the only Ai2 project, and being responsible with our resource use minimizes the disruption that large model training causes for other teams.

Below we describe several PyTorch optimizations 27 27 27 All of these techniques and more are implemented in the open-source OLMo-core library, the 2nd generation of the OLMo training codebase. OLMo-core is available at allenai/OLMo-core . that had a big impact towards reducing training time of LMs on our infrastructure without any apparent loss in the speed of convergence.

##### Taking advantage of compilation

torch.compile() is a function in PyTorch 28 28 28 pytorch.org/tutorials/intermediate/torch_compile_tutorial that will compile native PyTorch modules and functions into optimized kernels, resulting in significant throughput improvements and GPU memory savings by avoiding the Python overhead associated with calling individual PyTorch operations in sequence, and by reducing the number of reads and writes that must occur on the GPU. As such, when torch.compile() is used properly, it can effectively match the performance of hand-crafted kernels in many cases without the additional complexity and engineering effort ( Ansel et al., 2024 ) .

##### Minimizing host-device syncs

By default, GPU operations are asynchronous. A function call that uses the GPU is enqueued to a particular device, but not necessarily executed until later. This allows the system to execute more computations in parallel, including operations on CPU or other GPUs… – 113

Any time the training code forces a host-device sync , no more operations can be enqueued until all operations currently in the queue complete. These synchronization points will hinder performance, and it is easy to unintentionally introduce them.

A surprising number of operations cause host-device syncs: 1. Synchronously copying a tensor from CPU to GPU (e.g., with tensor.to(device="cuda") ) will force a host-device sync. This can be avoided by copying the tensor asynchronously (e.g., with tensor.to(device="cuda", non_blocking=True) ).

2. Copying a tensor from GPU to CPU cannot safely be done asynchronously, so GPU → \rightarrow CPU data transfer should be avoided whenever possible. Seemingly innocuous code can cause GPU → \rightarrow CPU data transfer, and therefore host-device syncs, such as print -ing a CUDA tensor or an “ if ...: " block that depends on how a CUDA tensor resolves to a boolean value.

3. Specific PyTorch operations like masked_select() may unexpectedly cause a host-device sync 29 29 29 github.com/pytorch/pytorch/issues/12461 .

Host-device syncs can be detected by calling torch.cuda.set_sync_debug_mode("warn") before starting the training loop. This will cause PyTorch to emit a warning whenever a host-device sync occurs. This happens on a best-effort basis. Some syncs may still be missed.

##### Asynchronous bookkeeping with a separate backend

A typical training loop involves periodic “bookkeeping” operations like logging metrics and saving checkpoints. While these operations may be relatively fast, their aggregate cost over the course of a training run can be significant. These operations also usually involve host-device syncs. For example, a training metric like cross-entropy loss is the result of computations that occur on the GPU, and it is materialized first as a CUDA tensor; therefore, logging that metric to the console forces a synchronization point.

Many of these operations are essential and cannot be avoided, but it is possible to minimize the time they spend blocking the training loop by performing most of this bookkeeping work asynchronously, in a separate thread. However, the PyTorch NCCL backend is not thread safe. To work around this problem, we set up a separate backend that does not rely on NCCL (like GLOO), and use it exclusively for bookkeeping operations. The bookkeeping workflows could then look like this:

1. For metric collection and logging: Decide on the interval in which to log metrics. Since this involves a host-device sync, it should not be done on every training step. More commonly, metrics are logged every 10 or every 50 steps. During every step, metrics are computed and stored in a GPU tensor on their original devices. Only when it is time to log metrics do we copy them to the CPU (causing a host-device sync), and then pass them to the bookkeeping thread, which uses its own PyTorch backend to aggregate the metrics and log them.

2. For checkpointing: A similar workflow can be used for checkpointing. When it is time to save a checkpoint, the trainer makes a copy of the model and optimizer state in CPU memory (causing a host-device sync). Then it passes the copy to the bookkeeping thread, which assembles the model from the model shards that are stored on each compute node, and saves it to disk, while the main thread can continue training 30 30 30 See pytorch.org/blog/reducing-checkpointing-times for a concrete example. .

In both cases, the only impact on training time is one host-device sync, and the time it takes to copy data from the GPU to the CPU. The frequency of each event can be configured, and the overall impact on training time is negligible.

##### Explicit Python garbage collection

During training, the default Python garbage collector periodically runs a collection. In a distributed setting, with thousands of training processes that are expected to run in lock-step with each other, nothing enforces that these garbage collections happen at the same time on every process. Since distributed training can only proceed as fast as the slowest process, this causes a noticeable decrease in average training time per step as well as an increase in variability (Figure 16 ). Both worsen as the number of processes increases.

To work around this problem, the OLMo 2 trainer disables automatic garbage collection (e.g., by calling gc.disable() 31 31 31 docs.python.org/3/library/gc#gc.disable ). Then, it runs garbage collection explicitly at regular intervals, triggered at the same time in each process (e.g. by calling gc.collect(1) 32 32 32 docs.python.org/3/library/gc\#gc.collect ).

Model Total GPU Power (MWh) Power Usage Effect. Carbon Intensity Carbon Emissions Water Usage Effect. Total Water Usage (kL) Llama 2 7B 74 1.1 - 31 1.29 - 4.26 105 - 347 Llama 3.1 8B 1,022 1.1 - 420 1.29 - 4.26 1,450 - 4,823 OLMo 7B 104 1.1 0.610 70 4.26 487 OLMo 2 7B 131 1.2 0.332 52 1.29 202 OLMo 2 13B 257 1.12 0.351 101 3.10 892

### 6.5 Environmental Impact

Following our analysis in Groeneveld et al. (2024) and previous literature ( Patterson et al., 2021 ; Dodge et al., 2022 ; Luccioni et al., 2022 ; Li et al., 2023a ) , we estimate the environmental impact of training our final models by first calculating the total energy consumed during pretraining, and multiplying it by the carbon intensity of the local grid to estimate the amount of carbon released. We additionally extend our previous analysis to also estimate water consumption, calculated by multiplying the power consumed by the water usage efficiency of both the power generation and the cooling hardware. As in Groeneveld et al. (2024) , we emphasize that while our reporting is standard practice, it does not account for other environmental impacts such as embodied emissions and water consumption of the hardware during manufacturing, transportation, and eventual disposal, and other lifetime operational impacts such as deployment and inference, and thus our estimates should be viewed as lower bounds. We report detailed results for our models in Table 19 .

As in Groeneveld et al. (2024) , we calculate the total power consumption for each model by measuring the power consumption of an individual node every 25ms, calculating the average consumption throughout training, and multiplying by the total number of nodes. We then multiply this quantity by the power usage effectiveness (PUE) factor for the data center we use to train a model to account for the overall energy efficiency of the data center. As the majority of training for OLMo 2 7B is done on the Jupiter cluster, we use Jupiter’s efficiency metrics for our analysis of the 7B model. OLMo 2 13B is trained on Augusta; therefore, we use its efficiency metrics instead. We estimate consumption at about 391 MWh of energy by pretraining OLMo 2 7B and 13B.

To calculate carbon emissions, we multiply the total power consumption by a carbon intensity factor based on the physical location of each data center, measured in kg CO 2 per kWh. The Jupiter cluster is powered by Austin Energy, which most recently reported a carbon intensity of 0.332 kg CO 2 per kWh. 33 33 33 austinenergy.com/-/media/project/websites/austinenergy/commercial/carbonemissionscalculator.pdf The Augusta cluster is located in Iowa, and the state of Iowa has an average carbon intensity of 0.352 kg CO 2 per kWh 34 34 34 www.eia.gov/electricity/state/iowa , which we use for our calculations. We estimate that training our latest models emitted about 154 tCO 2 eq . C ​ O 2 ​ Emissions = P GPU ⋅ PUE ⋅ Carbon Intensity CO_{2}\hskip 3.0pt\text{Emissions}=P_{\text{GPU}}\cdot\text{PUE}\cdot\text{Carbon Intensity} To calculate water consumption, we multiply the total power consumption by the water usage effectiveness (WUE) of both the offsite power generation as well as the onsite cooling hardware. Both clusters use highly efficient, closed-loop cooling hardware, so we assume a WUE onsite of 0 liters per kWh. Following Reig et al. (2020) , we assume a WUE offsite of 1.29 L per kWh for our Jupiter cluster and 3.10 L per kWh for our Augusta cluster. We estimate that training our latest models consumed about 1.1 million liters of water . Water Consumption = P GPU ⋅ PUE ⋅ ( WUE onsite + WUE offsite ) \text{Water Consumption}=P_{\text{GPU}}\cdot\text{PUE}\cdot(\text{WUE}_{\text{onsite}}+\text{WUE}_{\text{offsite}}) Though we aim to report a comprehensive analysis of the environmental impact of training our models, we emphasize that this is a lower bound on the total cost of developing large models. In an upcoming paper ( Morrison et al., 2025 ) , we will provide more comprehensive analysis covering energy, emissions, and water consumption throughout model development, pretraining, and deployment.

## Conclusion

We introduce OLMo 2 and OLMo 2-Instruct , a family of fully open 7B, 13B and 32B parameter language models trained on up to 6T tokens. Both the base and instruct models are competitive with other open-weight models in their size categories such as Qwen 2.5, Gemma 2, and Llama 3.1. We detail the substantial contributions required to build competitive language models—many of which are different from the original OLMo —including stable infrastructure, architecture improvements for stability, innovations in late-stage training data, the latest post-training techniques, and many more details. We release all training and evaluation code, datasets, checkpoints, and logs required to reproduce and expand on the models. OLMo 2 marks continued progress in open-source language models, building a new ecosystem for research, one where new training methods and techniques need to be understood and shared.

## Author Contributions

A successful team project like OLMo would not be possible without the fluid contributions of many teammates across formal team boundaries. As not all of these can be captured, we indicate each authors’ primary contributing role in OLMo 2 . Authors are listed in alphabetical order:

• For base model development, including training and data curation: Shane Arora, Akshita Bhagia, Christopher Clark, Allyson Ettinger, Dirk Groeneveld, Yuling Gu, David Heineman, Matt Jordan, Jiacheng Liu, Kyle Lo, William Merrill, Tyler Murray, Jake Poznanski, Dustin Schwenck, Luca Soldaini, Oyvind Tafjord, David Wadden, and Pete Walsh.

• For instruct model development, including training and data curation: Faeze Brahman, Pradeep Dasigi, Nouha Dziri, Yuling Gu, Shengyi Huang, Hamish Ivison, Nathan Lambert, Saumya Malik, Lester James V. Miranda, Jacob Morrison, Valentina Pyatkin, Oyvind Tafjord, and Christopher Wilhelm.

• For operational support, including program management, legal guidance, release process, and more: Taira Anderson, David Atkinson, Crystal Nam, and Aman Rangapur.

• For Ai2 cluster setup and support: Michal Guerquin, Michael Schmitz, Sam Skjonsberg, and Michael Wilson

• For mentorship and advising: Ali Farhadi, Hannaneh Hajishirzi, Pang Wei Koh, Noah A. Smith, and Luke Zettlemoyer.

Authorship for this work was determined by those making direct contributions to the OLMo 2 models, related artifacts, and their release. Core contributors are recognized for their sustained, significant contributions critical to the success of the OLMo 2 project.

## Acknowledgments

This work would not be possible without the support of our colleagues at Ai2: