Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations

Authors: alexander  hägele
elie  bakouch
atli  kosson
loubna ben allal
leandro  werra
martin  jaggi
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:09:16

Abstract: Scale has become a main ingredient in obtaining strong machine learning models. As a result, understanding a model's scaling properties is key to effectively designing both the right training setup as well as future generations of architectures. In this work, we argue that scale and training research has been needlessly complex due to reliance on the cosine schedule, which prevents training across different lengths for the same model size. We investigate the training behavior of a direct alternative --- constant learning rate and cooldowns --- and find that it scales predictably and reliably similar to cosine. Additionally, we show that stochastic weight averaging yields improved performance along the training trajectory, without additional training costs, across different scales. Importantly, with these findings we demonstrate that scaling experiments can be performed with significantly reduced compute and GPU hours by utilizing fewer but reusable training runs. Our code is available at https://github.com/epfml/schedules-and-scaling/.

Main paper: 
 introduction: 
 Training large language models is expensive -in time, energy, and compute. Moreover, it requires a complex algorithmic recipe of model architecture and training data to obtain high-quality models. Therefore, the workflow of training large models consists of iterating over small experiments to verify success before extrapolating to larger scales. This is then either done by computing specialized scaling laws (OpenAI, 2023; Bi et al., 2024; Hu et al., 2024; Team et al., 2023) , relying on established laws (Hoffmann et al., 2022; Anil et al., 2023) or training past compute-optimality to save cost at inference (Touvron et al., 2023a,b) .Despite large advances across data and training recipes, one aspect of large language model (LLM) pretraining has remained surprisingly prevalent: the cosine learning rate schedule (Loshchilov & Hutter, 2016; Radford et al., 2018; Rae et al., 2021) . Importantly, the Chinchilla project (Hoffmann et al., 2022) showed that the cosine schedule achieves optimal loss only when the cycle length matches the training duration, but underestimates the model performance during training. This means that when performing experiments -e.g., for architectural changes or data mixtures -one must train multiple models for different lengths, from scratch, to have reliable estimates of the quality of training and the scaling behavior. This is much more expensive than training a suite of models just once. Even more, it is restrictive for the final model for which the training length must be decided in advance.In this work, our goal is to revisit and question the necessity of the cosine learning rate schedule for large model training. Through a multitude of training runs, we demonstrate how a simple alternative of performing a cooldown after a constant learning rate -which was already suggested in the literature (Zhai et al., 2022) and recently used by released models (Hu et al., 2024; Shen et al., 2024) -matches the performance of cosine. We expand on this and provide and analyze different recipes for the decay form and length, which scale as reliable as cosine, outperforming it for sufficiently long optimality Figure 1 : Revisiting cosine optimality for language models. We revisit the observation from Chinchilla (Hoffmann et al., 2022) that in order to achieve the best model after a certain training length (tokens), the cosine schedule must match the total duration of training. This comes at the cost of neither being able to stop before or going beyond the cycle -an issue we show how to alleviate in Section 3. cooldowns. Going beyond, we investigate stochastic weight averaging (Izmailov et al., 2018) and a schedule-free optimizer (Defazio et al., 2024) , which give strong (but not optimal) performance at any point during training and can act as a replacement for the learning rate decay, if the performance gap is acceptable and avoiding separate cooldowns is preferred.These findings suggest that research on training recipes and scaling laws has been needlessly complex due to the need to retrain models from scratch. We demonstrate this empirically by performing a small-scale experiment of scaling laws which only uses a fraction of the compute and GPU hours that were previously needed. Following, we discuss that this makes scaling research more accessible and enables more frequent computation of laws for data mixtures (Bi et al., 2024; Goyal et al., 2024; Aghajanyan et al., 2023) or novel architectures (Gu & Dao, 2023; De et al., 2024) .2 Background: Cosine Learning Rate Schedule for LLMs Revisiting the optimality of the cosine schedule. We start our argument by revisiting the use of the cosine schedule in large language model (LLM) training. For any machine learning model, the learning rate value (LR) and schedule both are crucial choices for training. From optimization theory, our understanding is that a slow annealing of the learning rate is essential to find good minima in the loss landscape particularly for deep networks, whereas higher values help exploration (Smith et al., 2017; Loshchilov & Hutter, 2016) .In the context of LLMs, the most commonly used cosine strategy presents a particular trade-off by which the LR reaches its maximum early after the warm-up stage and then gradually decreases, typically to 10% of the maximum LR (see Figure 1 , right). Since the seminal works of GPT (Radford et al., 2018 (Radford et al., , 2019;; Brown et al., 2020) and large models like Gopher (Rae et al., 2021) , PaLM2 (Anil et al., 2023) or LLaMA (Touvron et al., 2023a,b) , cosine has stayed the de-facto standard schedule.Experimental setup. Our first goal is to understand the importance of the length of the schedule for performance of the model. To this end, we implement the common decoder-only transformer (Vaswani et al., 2017) identical to the LLaMa (Touvron et al., 2023a,b) or Noam architecture (Ormazabal et al., 2024) . Throughout this paper, we use the AdamW optimizer with weight decay (Kingma & Ba, 2014; Loshchilov & Hutter, 2017) with common LLM training parameters. We train on a subset of SlimPajama (Soboleva et al., 2023) with 6B tokens1 , a cleaned and deduplicated corpus for LLM pretraining, which we split into train and validation sequences and report validation loss (perplexity). We provide all the details in Appendix A.1.The pitfalls of cosine. In the results of Figure 1 , we see that the key parameter for the cosine schedule is the length of training: At specific step counts, the best perplexity is always achieved by the cosine schedule that matches the length. This is the main observation of Hoffmann et al. (2022) -in order to achieve the best model for a specific token count, the training duration must be known in advance and match the cosine length. However, this brings particular issues. First, cosine is suboptimal during training and underestimates the model's performance for the same token count. At the same time, cosine strongly complicates continuation of training. For example, one could easily be mistaken to extrapolate a loss curve of a cosine schedule beyond the end of the cycle. The improvement in loss precisely happens because of the LR decay; afterwards, the final learning rate will generally be too low to continue making large progress. In contrast, rewarming leads to spikes of which the training only slowly recovers, similarly reported in the continual learning literature (Ibrahim et al., 2024) . The alternative -constant + cooldown. The tradeoff of cosine can also be achieved by a different schedule. Here, the learning rate (LR) is kept constant for the majority of the training and only decreases in a short final phase, known as the cooldown (or decay/annealing phase). This schedule has previously been referred to as a trapezoidal (Zhai et al., 2022) , and later as the warmup-stable-decay schedule (WSD) (Hu et al., 2024) . To avoid overloading of terms (e.g., weight decay), we refer to this approach as constant LR + cooldown. 
a different: 
 The cooldown phase typically has the LR go to zero, mirroring the warmup phase, which gives rise to an overall trapezoidal shape (see Figure 2 ). Formally, we can define it asη(n) =      n Nwarmup • η max if n < N warmup η max if N warmup < n ≤ N -N decay f (n, N, N decay ) • η max if n > N -N decay (1)with the peak learning rate η max , the total steps N with warmup and cooldown steps N warmup and N decay , and a monotonically decreasing function f (n, N, N decay ) that handles the cooldown.The main advantage of the constant schedule is that it does not require one to specify the number of training steps in advance. This is particularly convenient for large runs, as the cooldown can be initiated at any time to observe model behavior and decide whether to stop. We want to highlight the following in more detail.Continual learning. The constant + cooldown schedule allows for continual learning by default.Here, the natural approach is to use checkpoints before the cooldown to continue training with a high LR; this avoids loss spikes and brittle training when rewarming the learning rate (cf. Figure 1 with cosine). Moreover, rewarming has been reported to hurt performance and introduce forgetting compared to single annealing training (Ibrahim et al., 2024) , though careful strategies can alleviate such issues. It remains an interesting question if a single cooldown schedule is absolutely optimal given a total compute budget for LLM training.Data mixtures. During the cooling phase, the data mixture can be changed (Hu et al., 2024) as a form of finetuning; beyond aiming for specific downstream tasks, this allows one to assess the quality of specific mixes, e.g., as recently done for Llama 3 (Dubey et al., 2024) . Although we focus on the same mixture, understanding the curriculum aspect of the separate phases is explored in the concurrent work of Blakeney et al. (2024) and remains open for further research.Scaling studies. Since a cooldown can also be initiated retrospectively from a checkpoint, the schedule allows to see the same model at different scales of compute. This enables much cheaper scaling studies than commonly done; we show this in depth in Section 5.We follow our experimental setup from the previous section (details in A.1) and train a 210M parameter model on SlimPajama with constant LR and the cooldown schedule defined in (1). Figure 3 : The difference in loss curves of cosine vs. constant learning rate with cooldown. The cooldown phase initiates a sharp decrease in loss (left) to match cosine; the training perplexity follows the same behavior (Fig. 15 ). We find the LR sensitivity (right) to be similar for both schedules, albeit less for cooldown, where the optimum lies slightly below at half of the optimal cosine maximum LR. Perhaps surprisingly, we find that a different decay phase in the functional form of (1-sqrt) can consistently outperform the standard linear decay, where both are better than the chosen (untuned) cosine for long lengths.we compare the same length of warmup and training, but replace the cosine decay after warmup with a constant LR and a cooldown for 20% of steps linearly going to zero. We additionally sweep the maximum learning rate for both approaches. In the results shown in Figure 3 , perhaps suprisingly, we observe an almost perfect match between the performance of the best cosine and cooldown schedule even for different training durations, all while exhibiting slightly less sensitivity to variations in the LR.Different cooldown schedules. We investigate various functions describing the cooldown shape, including cosine, square, and square root shapes in Appendix B.1. We identify a new function (1-sqrt) that outperforms the linear cooldown. This improvement is maintained in a smaller number of decay steps, different learning rates, and various timestamps (see Appendix B.1). In Figure 4 , we train a model for 200K steps (approximately 20B tokens), applying cooldown every 20K steps for 20%. The results indicate that the longer the training duration, the more linear cooldown is outperformed by the (1-sqrt) cooldown, which we define as:f (n, N, N decay ) = 1 - n-(N -Ndecay) Ndecay (1-sqrt)Takeaway 1: The constant LR + cooldown schedule offers significant convenience by not requiring the number of training steps to be specified in advance, and provides similar performance compared to a well tuned cosine schedule.Takeaway 2: We find a cooldown form (1-sqrt) that consistently performs better than linear decay.How long do you need to cooldown? To effectively utilize the schedule, it is essential to determine the optimal number of decay steps. Our study of the relative number of cooldown steps, as shown in Fig. 5 , reveals that the benefits of extended cooldown periods plateau at around 20%, which we select for the remaining experiments. Additionally, in Fig. 6 we demonstrate that using only 5% decay with the (1-sqrt) cooldown can nearly match the performance of the cosine schedule on a 20B token run (much beyond Chinchilla optimal). This is practically important, as we ideally keep the cooldown as short as possible but get strong performance. Takeaway 3: The constant learning rate with a short cooldown (< 20% of total training steps) can achieve the same final loss as cosine, and only outperforms for a longer fraction of cooldown steps.For long training runs, our results suggest that the cooldown length can be less than 20% to match cosine if it is enough in absolute steps -see Figure 6 .What happens during the cooldown? It is remarkable how the sudden drop in loss is consistent for both train and validation loss and aligns closely with the decay in LR (Figure 15 ). Hu et al. (2024) investigate the cooldown phase and find that the first-order directional derivative diminishes with each step, whereas the curvature of the loss function increases; they attribute this to proximity to a local optimum.We expand upon these results and aim to understand the optimization landscape around the trajectory of the cooldown. For that, we evaluate the loss along the straight line trajectory when moving between a checkpoint before and after cooldown, i.e., a linear interpolation of the weights of the two models. Perhaps surprisingly, we find that the smooth drop in loss also occurs for this interpolation, as visualized in Figure 7 . This aligns with the findings of Hu et al. (2024) . These results suggest that upon decaying the LR, the model immediately descends into a connected minimum of the loss.Takeaway 4: The cooldown phase is a smooth transition to a basin in the loss landscape.Decaying cosine to less than 10%. Throughout our experiments, we follow the common heuristic of setting the final LR of cosine to 10% of the maximum. We also provide experiments in which we ablate annealing to zero (or a very small value) in Fig. 22 in Appendix B.1. We find that, perhaps unsurprisingly, the influence of the final LR is non-negligible: decaying to lower values improves performance to match the best (1-sqrt) cooldown. However, when evaluating on downstream benchmarks (see next Section 3.3), 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. Figure 8 : The performances of both schedules also match for downstream tasks. We run a realistic setting of a 1B model trained for 100B tokens of FineWeb and establish matching performance of both schedules for common LLM benchmarks. Interestingly, there is a similar boost in performance with the cooldown. Detailed numbers over the course of training and 460B token runs are in Appx. B.5.We expand our ablations to realistic scenarios of 1B and 8B parameter models, with an additional focus on model evaluations. Our goal in this section is therefore twofold: first, we validate the dynamics of the schedule at practical scale; second, we establish the connection between the pure loss values and downstream benchmarks. Although tasks often scale reliably with loss (Du et al., 2024; Gadre et al., 2024) , this connection is crucial, as downstream abilities are ultimately the main metric of interest.1B with downstream evals. We train a 1B parameter model with cosine and a cooldown schedule with the high-quality FineWeb corpus (Penedo et al., 2024) . Following DeepSeek scaling laws (Bi et al., 2024) , we set an approximately optimal batch size and LR of 1.8M and 8e-4 for 100B and 460B token runs; evaluation is done throughout training on the most common benchmarks like MMLU (Hendrycks et al., 2021) . The setup is described in detail in Appx. A.2. With the aggregated results in Fig. 8 , we find a clear agreement of the downstream performance for both schedules. Interestingly, there is an uptick for some evaluations, similar to the drop in loss, with the start of the LR cooldown; see Fig. 29 . However, not all benchmarks follow this trend, which requires further research. The results for the much longer 460B run, where performance also matches, and detailed benchmark numbers are provided in Appx. B.5.8B. While a full 8B training run is beyond our limits, we provide a first investigation in Fig. 9 : In line with all our experiments, we see matching loss values and, promisingly, no instability at such a large scale, though the two runs are clearly short (20k steps). Moreover, instabilities arising from a high LR for longer training can be alleviated by methods like QK norm (Wortsman et al., 2024) . 4 Do We Even Need to Cooldown?In Section 3, we show that a constant LR with a short cooldown phase can replace the cosine decay.Ideally, however, we would not need a decay at all, but aim to obtain an optimal model at any point in training, even when prolonged. This would save even more computational resources and time. In this section, we investigate two potential approaches: weight averaging and a schedule-free optimizer.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
stochastic weight averaging (swa): 
 Motivation. While slow annealing of the learning rate can be essential to find good minima (Smith et al., 2017; Loshchilov & Hutter, 2016) , Sandler et al. (2023) show theoretical and empirical equivalence between stochastic weight averaging (SWA) and cooldowns in training vision models.There, intuitively and naturally averaging reduces noise and thus improves generalization (Wortsman et al., 2022) . Motivated by these results, we aim to answer the same question in LLM training: Can weight averaging replace the cooldown phase?Method. We opt for a form of SWA (Izmailov et al., 2018) that splits the training in fixed windows and averages within a window, allowing us to keep the average as a single additional copy of the model parameters. In our experiments, we set the window to h = 500 steps and save the window averages as checkpoints every h steps, which allows ad hoc evaluation of longer windows to see potential improvements akin to latest weight averaging (LAWA, Kaddour, 2022; Sanyal et al., 2023) . For all our experiments, we find that windows below or at 2500 steps (256M tokens) are optimal.We also experimented with an exponential moving average (EMA), which performed worse than SWA, and therefore we do not report EMA. The cooldown schedule outperforms SFO even when tuning momentum parameters. We find that SFO is sensitive to the choice of (β 1 , β 2 ) momentum parameters. It gives strong performance for well-tuned momentum, but falls short of cooldown.Experimental results. We evaluate SWA with the same 210M model and show the results in Fig. 10 for both a constant LR (left) and cosine (right). Notably, we find a significant performance boost for SWA on top of a constant LR. Yet, it does not reach the loss values of explicit cooldowns. On the other hand, in line with previous work (Kaddour, 2022; Sanyal et al., 2023; Andriushchenko et al., 2023) , we see a similar boost for SWA on top of cosine. This suggests that, regardless of schedule, SWA is a compelling approach to achieving strong models along the datascale axis that can serve as a replacement for models trained with fewer steps, if the performance gap is acceptable and one wishes to avoid cooldowns. This is particularly advantageous as it can arguably be done for free on top of existing well-tuned optimizers.Takeaway 5: Irrespective of the schedule and without additional overhead, SWA improves performance along the training trajectory and provides better models during training. While it does not match the cooldown, it reduces the gap without the need for a separate decay phase.Very recently, Defazio et al. (2024) introduced a schedule-free optimizer (SFO) which uses an interpolation between standard averaging and Polyak-Ruppert averaging, inspired by Nesterov's accelerated method (Nesterov, 1983) . As such, the optimizer does not require a decreasing learning rate schedule, making it relevant for continual training. We seek to investigate if it can outperform the presented cooldown schedule and provide a comparison in the same LLM setting.Results. We compare the results of a long 210M model run with cooldown vs. SFO with AdamW in Figure 11 . Although the optimizer does not require a learning rate schedule, the authors point out that training is more sensitive to the choice of the momentum parameters (β 1 , β 2 ) (which are not identical to the momentum in Adam). They note that the optimal parameters may depend on the length of training, making it not fully schedule-free. We observe this sensitivity in our experiments, where the choice of (0.9, 0.95) performs significantly worse and even increases the loss toward the end of training. For (β 1 = 0.95, β 2 = 0.99), SFO performs remarkably well. Nevertheless, both settings are matched or outperformed by the cooldown schedule, in particular when comparing the same momentum configuration. We did not perform any further hyperparameter tuning for either method. 
the implications for scaling law research: 
 Our investigation shows the efficacy of a constant LR for flexible training of LLMs by introducing a short cooldown phase at any point during training. In this section, we go beyond a single model, and focus on the reliability of the alternative LR schedule compared to cosine across model sizes and scales. Ultimately, we discuss the importance for the future of scaling law experiments.Importance of scaling laws. At their core, scaling laws aim to establish a functional form of a model's performance, commonly modelled as the loss L, as a function of the parameters N or training tokens D; for LLMs, this is usually expressed as the power lawL(N, D) = A N α + B D β + E ,where {A, α, B, β, E} are variables to be estimated (Kaplan et al., 2020) . Such scaling laws serve a multitude of critical purposes -from optimally spending a fixed amount of compute (FLOPs) for achieving the lowest loss, to trading off data-sources of different quality (Bi et al., 2024; Goyal et al., 2024) or modalities (Aghajanyan et al., 2023) , to comparing the efficiency of different architectures (Gu & Dao, 2023; De et al., 2024) .Crucially, Hoffmann et al. (2022) demonstrated that it is necessary to vary the number of training tokens for a fixed family of models and fully decay the LR (see also Kaddour et al., 2023) . At the time, their results suggested that LLMs were over-sized, leading to a substantial increase in data collection and training for much longer, beyond the Chinchilla optimal point (N, D) (Touvron et al., 2023a,b) .Why do the presented results matter for scaling? Following the results of Chinchilla, scaling laws require a family of models each trained from scratch with a cosine schedule that is fit to different training lengths (cf. Section 2). In contrast, the cooldown schedule as well as weight averaging allow a much cheaper alternative in two phases: first, a model sweep with a single sufficiently long training run for each model size in the family; then, using the model checkpoints to perform a cooldown or averaging. This reduces scaling law experiments to only the model scaling axis, effectively dividing the number of necessary training runs by one order of magnitude. At the same time, it allows for flexible continual training beyond any predetermined number of steps.Experimental setup. We mimic a small-scale experimental setup for scaling laws: We train a range of model sizes (33M-360M) across different token scales (0.3B-10B) on the same SlimPajama 6B dataset. For each model, we choose exactly three token counts (around the Chinchilla optimal ratio of D/N=20) in increments of 10, 20 and 30 tokens per parameter. With the cosine schedule, each model is trained from scratch three times. In contrast, for averaging, we train each model just once for the longest cycle and then use the averages at the same token count as cosine; for cooldown, we similarly take checkpoints along the constant LR trajectory and perform three annealing periods to match the token counts. We adjust the LR to be higher for smaller models. In line with the findings of Sect. 3 and for a fair comparison, we set the constant LR to be half the maximum LR for cosine and perform 20% (linear) cooldown steps. More details are given in Appendix A.1.Results. We show the validation loss envelopes as a function of the training FLOPs in Fig. 12 (left) and compare each obtained model (i.e., same parameter & tokens) with cosine and its alternatives (right). We find that the cooldown learning rate schedule scales reliably, much like cosine and SWA, to achieve optimal losses with the recipe we establish in previous sections. This is particularly visible in Fig. 12 (right), where each model's performance lies almost perfectly on the diagonal, which signals that cosine (y-axis) and cooldown (x-axis) reach the same loss after the same amount of tokens. A similar result is visible for SWA, though the reciprocal (y-offset) is negative, which agrees with our findings from Sect. 4.1 that the averaging does not fully close the gap to a LR cooldown.Compute savings. Crucially, the alternatives enable scaling laws for just a fraction of the cost. In Fig. 13a , we report FLOPs and GPU hours (real wall-clock time) for all model runs for both methods. They substantially reduce both compute and GPU hours. In our experiments, where we space the runs to use token ratios of 10, 20 and 30, it saves half the time and FLOPs, enabling scaling laws for only a fraction of the previous cost. We report the detailed savings for all models in Appendix B.2.Estimating savings for Chinchilla. We take our analysis further and estimate how much cheaper the Chinchilla model suite would have been if performed with 10% cooldowns after a single run for each model size using the model configurations as reported in Table A9 (Hoffmann et al., 2022) . Since the authors do not report exact training configurations, we consider a sequence length of 1024, a batch size of 0.5M tokens and token ratios M = D/N ∈ {10, 15, 20, 25} for each model. With this, we arrive at roughly 5.59 × 10 23 total FLOPs originally vs. 2.36 × 10 23 (Figure 13b ). This means that less than half the compute could have been used.Takeaway 6: Scaling experiments can be done with significantly reduced compute and GPU hours by utilizing fewer but reusable training runs with a constant learning rate and ad-hoc cooldowns.Additional results. We plot the training curves of all models in Appx. B.3. In addition, we show the same findings with experiments on OpenWebText2 in Appx. B.4. 
limitations: 
 We conduct our experiments on models of up to 8B parameters, with long training runs of a 1B model on multiple hundred tokens. The trends we find are consistent across all scales, but training behavior can be more brittle at modern scales and extremely long training (Wei et al., 2022; Tay et al., 2021) . However, instabilities arising from a high learning rate for a large part of training can be alleviated (Wortsman et al., 2024) . 
related work: 
 Cosine Schedules and Alternatives for Transformers. The cosine decay was originally introduced by Loshchilov & Hutter (2016) for cyclic schedules in vision tasks, where it is common practice to have stepwise or cyclic LRs to escape bad minima when training multiple epochs (Smith et al., 2017) . For language models where data is more vast, the cosine schedule with a single cycle is currently the de-facto standard for training, with few exceptions of T5 (Raffel et al., 2020) and PaLM1 (Chowdhery et al., 2023) that used a form of inverse square root. In line with our work, recently released models opt for alternatives such as stepwise schedules (Bi et al., 2024) or the presented constant + cooldown (Shen et al., 2024; Hu et al., 2024) . These alternatives were previously also explored for vision transformers by Zhai et al. (2022) , who find reciprocal square-root with cooldown to perform best, and in the context of continual learning (Ibrahim et al., 2024; Gupta et al., 2023) . Defazio et al. (2023) investigate the gap between theory and practice of LR schedules and suggest that a linear decay is optimal, also for LLM training. We similarly find that a long linear decay can slightly outperform cosine.Weight Averaging. Weight averaging over past iterates (Polyak & Juditsky, 1992) has long been known to be beneficial for convergence. Izmailov et al. (2018) introduce stochastic weight averaging for better generalization in deep learning models. Similarly, the exponential moving average is commonly used in vision (Morales-Brotons et al., 2024) . Importantly, Sandler et al. (2023) show the equivalence of WA with decaying learning rate schedules. In the context of LLM training, close to our work is Sanyal et al. (2023) which showed that a form of latest averaging (Kaddour, 2022) can be used to improve the performance of models early in training. However, they did not investigate the relation of weight averaging to compute optimality and its implications for scaling experiments.Scaling Law Experiments for Neural Language Models. Kaplan et al. (2020) were the first to establish scaling laws for language models. Important to our work, Hoffmann et al. ( 2022) revise these findings and demonstrate specific methods to establish scaling laws, notably training a family of models for different cosine lengths. The subsequent models like LLama1/2 (Touvron et al., 2023a,b) further improve performance of smaller models by training beyond the Chinchilla optimal point, motivated by lower inference costs (Gadre et al., 2024; De Vries, 2023; Sardana & Frankle, 2023) . Recent works (Muennighoff et al., 2023; Bi et al., 2024; Goyal et al., 2024) highlight how data repetition and quality affect the scaling behavior, which suggests that scaling laws should be updated more frequently. However, these works do not consider efficient experiments for scaling laws, which is the focus of our work.In concurrent work, Porian et al. (2024) and Pearce & Song (2024) find that the discrepancy between the Kaplan and Chinchilla scaling laws is not attributed to the learning rate decay, but they establish that the optimal token/parameter ratio can be obtained with a constant learning rate without any cooldown; however, the actual performance is then suboptimal, and a cooldown schedule as suggested in our work is needed to properly estimate model performance, in particular for downstream tasks. 
conclusion: 
 We have demonstrated the reliability of an alternative learning rate schedule to replace cosine for LLM training, which uses a constant rate with a cooldown phase. Across a multitude of experiments, we analyze different recipes for the decay form and length. Importantly, we do not claim to have established the best learning rate schedule -instead, we investigate and demonstrate how an arguably simple recipe can match the performance of the current best practice of cosine, and discuss how it provides compelling advantages such as continual training and a strong reduction in costs for scaling law research. In addition, we find that SWA can give reliable (strong, but not optimal) estimates of models during runs, without additional overhead or training.We believe the results are of great importance to the present and future of LLM training: The presented methods facilitate research for the current post-Chinchilla era, where models are trained much beyond compute-optimal, by allowing more flexibility to continue training whenever needed. At the same time, recent results that suggest data dependency in scaling (Bi et al., 2024; Goyal et al., 2024; Aghajanyan et al., 2023; Pandey, 2024) imply the need to frequently update scaling laws, which is economically more feasible with reduced costs. We therefore hope that our work will make scaling research more accessible to researchers and practitioners alike. and gradient clipping with 1.0. For the 1B model, we use a custom BPE tokenizer5 (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 here6 7 . We report the accuracy normalized by sequence length (acc_norm). Table 3 : Model configurations for larger runs. The batch size and learning rate for the 1B model tokens were estimated using DeepSeek scaling laws for 100B tokens. The two values for BS and the total steps of the 1B model distinguish the runs (100B,460B). For the 8B model, the architecture is identical to Llama3 and the batch size was set according to the available GPU limits on our cluster at the time of running experiments. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not release pretrained models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: To the best of our knowledge, we adequately cite and mention all used assets in this paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: We do not provide new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: We do not involve crowdsourcing or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: We do not involve crowdsourcing or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: 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 .1 def embedding ( seq_len , vocab_size , d_model ) : 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 ) where a = 0.5, we sweep the parameter 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, 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.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 . The savings amount to a factor of 1 2 across all models. If more runs across different lengths were to be performed (compared to 3 per model in our experiments), the difference would be even more significant. 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 .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 We train a 1B model on 100B (left) and 460B tokens (right) of FineWeb (Penedo et al., 2024) , and find that the performance of cosine and cooldown matches. Though cosine to zero improves the loss (Figure 22 ), it leads to a saturation before the end of training, hurting overall performance.Figure 27 , our findings from previous experiments succesfully transfer, where the cooldown recipe matches the performance of cosine and SWA boosts performance during training.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 4 (100B) and Justification: We explicitly discuss limitations in Section 6. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.https://huggingface.co/datasets/DKYoon/SlimPajama-6Bhttps://x.com/BlancheMinerva/status/1754559681933324372https://github.com/karpathy/nanoGPThttps://github.com/huggingface/nanotronhttps://hf.co/lvwerra/the-tokenizer-v1https://github.com/huggingface/cosmopedia/blob/d62abb8e13c567b999e33d0a1d795968bc052c6a/ evaluation/README.mdhttps://hf.co/datasets/HuggingFaceFW/fineweb/blob/main/lighteval_tasks.py#L12 


Decision: The paper offers a practical and well-validated approach to optimizing training schedules for LLMs, with significant implications for continual learning and reducing computational costs in scaling law studies. While the technical novelty is modest, the thorough experimental analysis and potential impact make this work a strong contribution to the field. All reviewers are in favor of this paper. This is a clear accept.

Reviews: review: paper_summary: paper_summary: The paper focuses on learning rate schedules in scaling large language models (LLMs). Traditionally, LLMs are trained with a cosine learning rate schedule, which requires a pre-specified training duration for learning rate decay. This makes it difficult to dynamically adjust the training duration, as early stopping and continued training after the decay result in suboptimal model performance. To address this, the authors study an alternative schedule of constant learning rate with cooldown, which enables one to get the optimal model at any step during the training process. The authors systematically compare different choices of cool down duration and cool down schedule, and demonstrate that constant learning rate with 1 - sqrt cooldown schedule matches the performance of cosine schedule tuned for that duration.

Beyond cooldown, the authors compare two different approaches of removing the need of learning rate schedule: stochastic weight averaging and schedule-free optimizer. The authors show that while they exhibit promising performance, the performance of models obtained in the middle of the training still cannot match the performance of the proposed constant learning rate with cooldown.

Finally, the authors validate the proposed learning rate schedule in the scaling law study of compute optimal models, and demonstrate that the constant learning rate with cooldown schedule leads to the same results as cosine learning rate schedule, while saving half of the compute flops.
strength_weakness: strength_weakness: While the experiment results are very solid, I do hope that the authors could expand the scale of experiments to larger models. Larger models are usually more sensitive to suboptimal hyperparameters, and therefore it would be great if the proposed method can also be verified in larger models, such as a 3B or 7B model which are common sizes for practical LLMs.
questions: questions: Another highly relevant use of training LLMs without fixed duration is continued pre-training, which is quite common in practice as people often continue the pre-training on specialized data for different use cases (such as turning general models into code specialized models). It would be great if the authors could include some discussions and experiments on how the proposed constant learning rate with cooldown can be applied in continued pre-training. For example, one question to ask would be would it be better to continue from the checkpoint before or after cooldown.
limitations: limitations: The authors have sufficiently addressed the limitations and societal impacts.
paper_summary: paper_summary: Scaling Laws and Compute-Optimal Training without Fixed Training Duration

Summary
A major weakness of the cosine annealing learning rate schedule, one of the most prevalent learning rate schedules in LLM training, is that for optimal performance the cycle length must be adjusted based on the training duration. Intermediate checkpoints generally do not achieve optimal loss. This is inconvenient since achieving the optimal loss for various training durations requires individual runs for each duration. The authors investigative alternative scheduling methods for alleviating this difficulty. Namely, the authors seek solutions which allow provide near optimal intermediate checkpoints along one long training run, potentially with a relatively small overhead from applying cooldown or averaging. Such solutions are especially useful for computing scaling laws and hyperparameter tuning by reducing the computation cost significantly. Through careful experimentation the authors provide several insights about the following solutions (which have been proposed previously in the literature).

1.	Trapezoidal LR schedule with warmup + constant lr + cooldown
2.	Stochastic Weight Averaging
3.	Schedule Free Optimizer from Defazio et al. [1]

In particular it is demonstrated that the trapezoidal schedule is a simple and effective solution for provide near optimal checkpoints along a potentially infinitely long trajectory. 

[1] Defazio, A., Mehta, H., Mishchenko, K., Khaled, A., Cutkosky, A., & others. (2024). The road less scheduled. arXiv preprint arXiv:2405.15682.
strength_weakness: strength_weakness: The technical novelty is a bit limited because all the approaches evaluated have been proposed before in the literature (modulo some small tweaks). However, I do not view this as a major weakness since a thorough, well-presented review of these approaches is still a strong contribution.
questions: questions: 1. In figure 1 why are there no spikes? For example the n=8k curve is rewarmed every 8k iterations. Are the curves smoothed?
2. Can the SWA be clarified a bit? Is every evaluation done on the iterate average over all checkpoints in the last 500 steps? How many checkpoints does that include?
3. What about a cooldown function like (1-x^a) where a < 0.5? All the other cooldown functions tried seem to lie "above" 1 - sqrt(x).
limitations: limitations: Yes.
paper_summary: paper_summary: The authors present an empirical study on using constant learning rates (plus a short cooldown) instead of cosine schedules for training LLMs. The authors show that:

- Constant LR + cooldown roughly matches cosine schedules. (Fig 3,4)
- SWA of a long schedule almost matches the performance of shorter schedule given the same number of steps. (Fig 8).
- Chinchilla-type scaling laws can be derived with constant LR schedules + cooldowns which will save compute. (Fig 10)
strength_weakness: strength_weakness: - The experimental setup is unclear. But this should be very easy to fix in an updated version.
- It is still unclear to me exactly how well the match between constant LR and cosine is, especially for longer training. 
- The results on SWA and constant LR seems almost orthogonal. 
- Higher performance is not reached, so the only upside is in doing scaling studies. This might only be relevant for a few well-funded labs.
questions: questions: 1. Why plot perplexity instead of val loss?
2. Could you have a subsection called experimental setup? On line 104 you mention “. We follow our experimental setup and train a 210M model”, but I don’t think the details of this model is given earlier.
3. Based upon Figure 2, it seems like your schedule goes to 0 whereas cosine typically does not go to 0. Is this a confounder in your experiments?
4. Fig 4 shows constant LR beating cosine while Fig 3 shows them matching. Why this discrepancy?
5. Fig 6 (right) seems to show that there is still a small gap between constant (1-sqrt) and cosine – is this correct?
limitations: limitations: na

Comments: title: Re: Rebuttal
comment: Thank you for your reply and nice words!

> Maybe I'm missing something but I don't see any loss jumps in the left panel of Figure 1? All the curves are smoothly decreasing.

We think the submitted version of Figure 1 should contain the black dashed line for a cycled schedule ('Cycle n=8k'), which shows a loss jump at steps 8k and 16k -- is the Figure different for you? The line is partially hidden under the very light red one for n=8k, which might lead to the confusion.
Thanks for the replies!

1. Maybe I'm missing something but I don't see any loss jumps in the left panel of Figure 1? All the curves are smoothly decreasing.
2. Got it, thanks for clarifying.
3. Great, looking forward to seeing what comes out.

After reading the discussion with other reviewers and author responses I believe the final paper will be very nice so I will increase my confidence and maintain the high score.
I'd like to thank the authors for answering my questions and providing us with additional experiments. My concerns and questions have been fully addressed.