Source-linked AI summary
Continual Pre-Training of Large Language Models: How to (re)warm your model?
Kshitij Gupta, Benjamin Thérien, Adam Ibrahim, Mats L. Richter, Quentin Anthony, Eugene Belilovsky, Irina Rish, Timothée Lesort
TL;DR
Continual pre-training could update language models on new data without retraining from scratch, but distribution shifts can degrade performance on past data. This paper studies learning-rate re-warming on Pythia 410M models continued from the Pile to SlimPajama and finds improved downstream performance relative to training from scratch.
Problem
Continual pre-training seeks to update models with new data instead of retraining, while maintaining performance on previous data despite catastrophic forgetting and learning-rate decay.
Method
The study varies warmup strategies, maximum learning rates, checkpoints, and warmup lengths while continuing Pythia 410M pre-training from the Pile to SlimPajama.
Results
Re-warmed models improve performance over models trained from scratch, while higher maximum learning rates favor SlimPajama adaptation and smaller rates preserve more Pile performance.
Takeaways & Limitations
The results motivate continual pre-training on new datasets rather than restarting training from scratch.
Takeaways & Limitations
The conclusions may not generalize beyond the high-similarity, overlapping-data setup or the 410M-model and 297B-token scales, and repeated updating remains unverified.
Abstract
from arXiv · showhide
Large language models (LLMs) are routinely pre-trained on billions of tokens, only to restart the process over again once new data becomes available. A much cheaper and more efficient solution would be to enable the continual pre-training of these models, i.e. updating pre-trained models with new data instead of re-training them from scratch. However, the distribution shift induced by novel data typically results in degraded performance on past data. Taking a step towards efficient continual pre-training, in this work, we examine the effect of different warm-up strategies. Our hypothesis is that the learning rate must be re-increased to improve compute efficiency when training on a new dataset. We study the warmup phase of models pre-trained on the Pile (upstream data, 300B tokens) as we continue to pre-train on SlimPajama (downstream data, 297B tokens), following a linear warmup and cosine decay schedule. We conduct all experiments on the Pythia 410M language model architecture and evaluate performance through validation perplexity. We experiment with different pre-training checkpoints, various maximum learning rates, and various warmup lengths. Our results show that while rewarming models first increases the loss on upstream and downstream data, in the longer run it improves the downstream performance, outperforming models trained from scratch$\unicode{x2013}$even for a large downstream dataset.
1. Introduction
The paper studies continual pre-training as a cheaper alternative to retraining, focusing on whether re-warming the learning rate improves adaptation to new data while limiting forgetting. Experiments show that maximum learning-rate choices trade off upstream retention and downstream adaptation.
- Motivation: Continual pre-training updates existing models on new data to avoid complete retraining, while seeking low loss on both new and previous data.The approach faces catastrophic forgetting and eventual learning-rate decay across repeated training stages.
- Approach: The study re-increases the learning rate for Pythia 410M models moving from the Pile to SlimPajama, varying checkpoints, maximum learning rates, and warmup lengths.The setup tests whether re-warming can improve learning efficiency without retraining from scratch.
- Findings: Progressive warmup is unnecessary, whereas starting directly at the maximum learning rate causes an initial stability gap without later consequences.The initial loss spike is described as a chaotic phase that does not affect later performance.
- Findings: Higher maximum learning rates strengthen SlimPajama adaptation, while smaller rates preserve more Pile performance.Maximum learning rate therefore controls the tradeoff between downstream adaptation and upstream retention.
- Findings: Continual pre-training with the latest pre-trained checkpoint improves performance compared with training from scratch.This result is reported even though the downstream dataset is large.
2. Setup
The experiments continue pre-training Pythia 410M models from the Pile on a large SlimPajama dataset that overlaps in source data. They vary learning-rate schedules and checkpoints to study continual adaptation without replay.
- Datasets: The Pile serves as the upstream pre-training dataset, while SlimPajama serves as the downstream dataset for continual pre-training.SlimPajama is an extensively deduplicated version of RedPajama built from the LLaMA dataset.
- Datasets: SlimPajama contains approximately 297B training tokens and is sampled from sources similar to the Pile, allowing some upstream data to recur.The setup is analogous to training first on part of a dataset and later on all samples.
- Model: The model is the 410M-parameter Pythia architecture pre-trained on the Pile.The experiments use the GPT-NeoX model family.
- Optimization: Experiments vary maximum learning rates of 1.5 · 10^-4, 3 · 10^-4, and 6 · 10^-4 under cosine decay to 0.1 · MaxLr.Warmup lengths are computed relative to the 297B-token downstream dataset, with the schedule reaching its minimum at 240B tokens.
3. Related Work
Prior work addresses continual learning, language-model domain adaptation, and learning-rate monitoring, but does not specifically examine warm-up phases for continual pre-training of large language models.
- Learning-rate schedules: Large language models commonly use a warm-up followed by cosine decay to 10% of the maximum learning rate.This pattern appears across models trained with Adam or AdamW.
- Continual learning: Continual pre-training has been studied in vision and language settings, including domain adaptation and datasets generated at different times.The paper’s setup is closer to domain adaptation because it does not model data temporality.
- Learning-rate monitoring: Prior continual-learning work evaluates constant, monitored, and otherwise adjusted learning-rate strategies because sequential datasets can reduce plasticity or cause forgetting.The paper places its warm-up study within this broader learning-rate literature.
- Research gap: No existing work is identified as specifically studying the influence of the warm-up phase in continual pre-training for large language models.This defines the paper’s stated research gap.
4. Continual Warm-up
The experiments examine how warm-up length, maximum learning rate, and checkpoint choice affect continual pre-training on SlimPajama while retaining performance on the Pile. Rewarming improves long-run downstream adaptation but creates trade-offs with upstream retention and does not benefit from longer warm-up or earlier checkpoints.
- How long to warm up?: Warm-up length does not significantly affect validation perplexity on either SlimPajama or the Pile.Experiments compare 0%, 0.5%, 1%, and 2% warm-up over a 297B-token schedule.
- How long to warm up?: Removing progressive warm-up causes an initial loss spike, or stability gap, without later consequences.The spike occurs during the first few training iterations.
- How high to warm up?: Larger maximum learning rates improve final downstream performance but worsen upstream performance through increased forgetting.Smaller learning rates preserve more Pile performance while limiting adaptation to SlimPajama.
- How high to warm up?: Constant learning rate training is initially advantageous, but its benefit disappears with sufficiently long SlimPajama training.Early stopping can retain strong upstream performance, whereas rewarming and decay performs better downstream at convergence.
- Comparing with from Scratch Training: Warm-started models outperform from-scratch training at convergence, including when the downstream dataset is large and overlaps with upstream data.This result indicates positive transfer between Pile pre-training and SlimPajama continual pre-training.
- Re-warming on the same data: Rewarming on the Pile produces a similar downstream-loss pattern, indicating that optimization dynamics contribute to the initial degradation beyond distribution shift alone.Earlier checkpoints also fail to improve downstream performance and instead decrease performance on both datasets.
5. Discussion / Limitation
The study’s conclusions are constrained by the similarity and overlap between upstream and downstream data, as well as by its limited model and dataset scales and evaluation metrics.
- Data similarity and overlapping: Data overlap between the Pile and SlimPajama may limit how well the results generalize to different distribution shifts.The authors specifically note that domain-adaptation settings may produce different outcomes.
- Data similarity and overlapping: Even slight distribution shifts can significantly perturb learning dynamics, so results differ between fine-tuning on the Pile and SlimPajama.The paper connects this possibility to prior evidence from image classification.
- Scope and evaluation: The experiments use 410M-parameter models and a 297B-token fine-tuning dataset, leaving larger scales untested.The authors propose testing 3B- and 7B-parameter models and datasets ranging from 100B to 600B tokens.
- Scope and evaluation: The study evaluates loss and perplexity rather than capability benchmarks such as HELM or Harness.The authors identify broader benchmarking as future work.
6. Conclusion
The experiments support continual pre-training with rewarming rather than restarting from scratch, while showing a trade-off between downstream adaptation and upstream retention. The authors caution that broader validation is still needed.
- Conclusion: Higher maximum learning rates improve adaptation to SlimPajama, whereas smaller rates better preserve Pile performance.Both settings outperform models trained from scratch in the reported experiments.
- Conclusion: Rewarmed models outperform models trained from scratch even when the downstream dataset is large.The conclusion motivates continuing pre-training on new datasets instead of restarting training.
- Conclusion: Further research is needed across larger model scales, different distribution shifts, and repeated model updates.These boundaries are stated as open questions for validating the strategy.
Software and Data
The experiments rely on GPT-NeoX, DeepSpeed, NCCL, Apex, PyTorch, and the HuggingFace Transformers library.
- Software: The implementation uses GPT-NeoX, DeepSpeed, NCCL, Apex, PyTorch, and HuggingFace Transformers.These are the software components listed for the experiments.
A. Upstream loss when fine-tuning various checkpoints.
The checkpoint experiments compare upstream validation loss across fully converged, saturation-point, and halfway-to-saturation checkpoints, while another figure summarizes faster validation-loss reduction from pre-trained checkpoints than from scratch.
- Checkpoint comparison: Figure 8 reports Pile validation loss for fully converged, upstream saturation-point, and half-saturation checkpoints.The experiments are described in Section 4.5.
- Checkpoint comparison: Figure 9 states that starting from a pre-trained checkpoint reaches lower Pile and SlimPajama validation loss faster than training from scratch.The comparison covers both upstream and downstream validation loss.