Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations — part 3/3

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

## Acknowledgements

We thank Anastasia Koloskova, Amirkeivan Mohtashami and Thomas Wolf for helpful discussions regarding the paper and its implications. We also thank Stella Biderman, whose open discussions 2 2 2 https://x.com/BlancheMinerva/status/1754559681933324372 inspired parts of our experiments.

## Appendix A Experimental Details

### A.1 Overview

Architecture and training parameters. We implement the most commonly used decoder-only architecture with SwiGLU activations ( Shazeer, 2020 ) , RoPE embeddings ( Su et al., 2024 ) , RMSNorm ( Zhang & Sennrich, 2019 ) and alternating attention and MLP blocks. Unless otherwise noted, we follow standard practices in LLM training and use the AdamW optimizer with beta parameters ( β 1 , β 2 ) = ( 0.9 , 0.95 ) (\beta_{1},\beta_{2})=(0.9,0.95) , decoupled weight decay of 0.1 0.1 ( Kingma & Ba, 2014 ; Loshchilov & Hutter, 2017 ) and gradient clipping with 1.0 1.0 . For warmup steps, we use a short warmup of 300 300 steps for the majority of runs and 1000 − 3000 1000-3000 for longer runs (above 100k total steps). The cosine schedule decays the learning rate to 10% of the maximum learning rate. For most of our experiments, we use a batch size of 200, i.e., roughly 0.1M tokens for a sequence length of 512. The vocabulary is based on the GPT-2 tokenizer ( Radford et al., 2019 ) and contains 50304 tokens.

Dataset and evaluation. Our main body focuses on results using SlimPajama ( Soboleva et al., 2023 ) , a cleaned and deduplicated corpus that includes webcrawl, code, papers and other sources, which is commonly used for pretraining LLMs. We use a subset of the full corpus that comprises roughly 6B tokens and randomly sample a validation set of roughly 3M tokens. During training, we evaluate the models with a fixed set of 32 batches of sequence length 512 (the same context length as training) to establish validation loss curves. At the end of training, we compute the full validation set perplexity. We perform further experiments that verify our main findings with OpenWebText2 ( Gao et al., 2020 ) in Appx. B.4 .

Implementation and infrastructure. Our code is based on an extension of NanoGPT 3 3 3 https://github.com/karpathy/nanoGPT and uses PyTorch ( Paszke et al., 2017 ) as well as FlashAttention ( Dao et al., 2022 ) . We incorporate bfloat16 for memory and throughput, trained with mixed precision float32 parameters and bfloat16 activations ( Micikevicius et al., 2017 ) . All experiments (aside from 1B and 8B, see A.2 ) were performed using a cluster of A100 GPUs (both 40GB/80GB RAM) with 2 data-parallel (i.e. 2 GPUs per run). Some selected runs used a single node of 8 H100s. We estimate that the full cost of all experiments for this project (including prototyping) to amount to roughly 2500-3000 GPU hours.

Model configurations. We provide an overview of the model sizes and configurations in Table 2 and the parameters for training and length in Table 3 . All models for the scaling law experiments are trained with 300 warmup steps and a sequence length of 512.

### A.2 Setup for the 1B and 8B Runs

Setup. In order to efficiently scale to larger models, we use the nanotron library 4 4 4 https://github.com/huggingface/nanotron that enables all forms of training parallelism while being performant and flexible. The configuration and hyperparameters of the two models are described in Table 4 . Similarly to our other experiments, we use the most common setup of the AdamW optimizer with ( β 1 , β 2 ) = ( 0.9 , 0.95 ) (\beta_{1},\beta_{2})=(0.9,0.95) , weight decay of 0.1 0.1 , and gradient clipping with 1.0 1.0 . For the 1B model, we use a custom BPE tokenizer 5 5 5 https://hf.co/lvwerra/the-tokenizer-v1 (that is very similar to GPT-2) because the FineWeb dataset ( Penedo et al., 2024 ) was already pre-tokenized on our cluster. Each run for the 1B model was performed on 4xH100s GPUs. For the 8B model, we reuse the original Llama3 tokenizer ( Dubey et al., 2024 ) for the similar reason that FineWeb-Edu was pre-tokenized on our cluster; here, we use 12 nodes, each composed of 4xGH200 GPUs. The model is split within one node with tensor parallel (TP) set to 4.

Evaluation. For evaluation, we use the lighteval library ( Fourrier et al., 2023 ) that works directly with nanotron checkpoints. We run the most common LLM benchmarks including MMLU ( Hendrycks et al., 2021 ) , ARC ( Clark et al., 2018 ) , OpenBookQA ( Mihaylov et al., 2018 ) , PIQA ( Bisk et al., 2020 ) , HellaSwag ( Zellers et al., 2019 ) , CommonSenseQA ( Talmor et al., 2019 ) , SIQA ( Sap et al., 2019 ) , Winogrande ( Sakaguchi et al., 2021 ) , truncating large benchmarks to 1000 samples so that we could efficiently evaluate over the course of training. This is a custom setup that is equal to that used for FineWeb ablations; the setup to reproduce the eval is described here 6 6 6 https://github.com/huggingface/cosmopedia/blob/d62abb8e13c567b999e33d0a1d795968bc052c6a/evaluation/README.md 7 7 7 https://hf.co/datasets/HuggingFaceFW/fineweb/blob/main/lighteval_tasks.py#L12 . We report the accuracy normalized by sequence length ( acc_norm ).

### A.3 FLOPs computation

We do not rely on the heuristic of 6=ND for computing the FLOPs of a Transformer model, but use a more thorough computation that involves the embeddings, attention and MLP operations directly. For reproducibility and other researchers to use, we provide our Python code in Figure 14 .

## Appendix B Additional Results

### B.1 More Results on Cooldown

Ablation of functional form of cooldown. We provide additional experiments to show that our results are consistent and robust to different changes. For simplicity, we only cooldown to 50K steps here, but we found these results hold regardless of the step at which we decay. In Fig. 16 and Fig. 17 , we test different functions for the cooldown phase: • Linear: Classical linear decay.

• 1 - Sqrt: Defined in Eq. ( 1-sqrt ).

• Cosine: Similar to the classical cosine decay, but applied only during the decay phase.

• Mirror Cosine: The symmetric counterpart of the cosine function with respect to the linear decay.

• 1 - Square: The square root function in Eq. ( 1-sqrt ) is replaced by a square function.

Note that the order might change for substantially different cooldown lengths; we focus on 10% and 20% because they are practically relevant.

In addition, we perform a comparison of different exponents for the square-root function: since the (1-Sqrt) can be expressed as ( 1 − x a ) (1-x^{a}) where a = 0.5 a=0.5 , we sweep the parameter a < 0.5 a<0.5 . The results are shown in Figure 18 . Apart from 0.1 and 0.2, which perform noticeably worse because the learning rate is too low for many steps, the other exponents only show a marginal difference to a = 0.5 a=0.5 , which still comes out on top.

We aim to further investigate the impact of the functional form of the cooldown in future work.

Length of cooldown. We repeat our ablation from Sect. 3 with a 210M model in Figure 19 (fractional x-axis) and with the absolute number of decay steps (not fractional) in Figure 20 .

LR sensitivity. The optimal learning rate also transfers to different cooldown lengths in our experiments, see the results in Figure 21 .

Annealing cosine to less than 10%. We ablate the choice of the final learning rate for the cosine annealing in Fig. 22 , where we find that the final value should be set lower than simply 10% of the maximum. However, when evaluating on downstream benchmarks (see Section B.5 ), we see that annealing to zero can hurt metrics; we posit this comes from too early saturation. Combining these two arguments implies that the maximum and final LR should be set (and ideally swept over) independently.

### B.2 Additional Results and Compute Savings

We give the savings for all models used in the scaling experiments in terms of FLOPs in Figure 24 and GPU hours in Figure 25 .

### B.3 Learning Curves for All Models of Scaling Experiments

We provide the learning curves for all models in the scaling experiments in Figure 26 . The final validation perplexity for all models and methods is given in Figure 23 .

### B.4 Additional Experiments on OpenWebText2

In addition to all results on SlimPajama, we perform experiments on the commonly used benchmark of OpenWebText2 ( Gao et al., 2020 ) with models of sizes 60M, 93M and 166M. As shown in Figure 27 , our findings from previous experiments succesfully transfer, where the cooldown recipe matches the performance of cosine and SWA boosts performance during training.

### B.5 Full Results of Large Model Runs

In this section, we report the detailed results of the 1B runs with the downstream benchmarks.

In Figure 28 , we compare the aggregate metrics throughout training for both the 100B (left) and the 460B token run (right). We equally plot individual benchmark curves in Figure 29 for the 100B run with a zoomed view after 80B tokens in Figure 30 ; interestingly, we observe a similar uptick in performance for some metrics (e.g., MMLU, HellaSwag) with the cooldown, while others do not benefit as clearly (e.g., OpenBookQA). This observation is an interesting direction for further research.

The final numbers are given in Table 5 (100B) and Table 6 (460B). Notably, longer cooldowns do not necessarily improve the metrics (e.g. going from 5% to 20%).

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .