Learning Rate Warmup › Warmup in Multi-Stage and Continued Training

Modern large language model development rarely follows a single training run from scratch to deployment. Instead, it involves multiple training stages: pretraining on a large general corpus, continued pretraining on domain-specific data, supervised fine-tuning, and reinforcement learning from human feedback. Each stage transition creates a new set of challenges for learning rate scheduling, and warmup plays a role at every stage.

When you transition from one training stage to the next, the model's weights are in a good state (they were learned during the previous stage), but the optimizer's accumulated statistics may not match the new training distribution. The Adam moment vectors $m_t$ and $v_t$ reflect the gradient distribution of the previous stage. If the new stage uses a different dataset, a different loss function, or a different mix of tasks, the pre-accumulated moment estimates can be misleading.

The conservative approach is to reset the optimizer state entirely at each stage transition and apply a fresh warmup from a low learning rate. This ensures the optimizer re-adapts to the new gradient distribution before taking large steps. The cost is that you lose the momentum information from the previous stage, which means the optimizer may converge slightly more slowly at the start of each new stage.

A less conservative approach is to keep the optimizer state but apply a short warmup of a few hundred steps. This treats the residual moment estimates as prior information that can be updated quickly by the new stage's gradients. In practice, this approach works well when the stage transition is not too dramatic (for example, continuing pretraining on a different text domain) but may fail if the distribution shift is large (for example, switching from masked language modeling to reinforcement learning from human feedback).

One approach that has been studied for multi-stage training is the cosine annealing with warm restarts (SGDR) strategy. In this approach, training alternates between a warmup ramp and a cosine decay, cycling repeatedly. Each cycle starts fresh from the peak learning rate, performs a short warmup, then decays back to a minimum. The cycle length is often increased geometrically (each restart cycle is longer than the previous).

The motivation is that each warmup-decay cycle gives the model a chance to explore a new region of the loss landscape (during the high learning rate phase) and then refine its position within that region (during the decay phase). The repeated restarts act like a sequence of increasingly refined searches. In practice, this strategy works well for continued pretraining and fine-tuning scenarios where you want the model to adapt to new data without forgetting the structure learned in previous cycles.

When fine-tuning a pretrained language model on a downstream task, the warmup requirements are quite different from pretraining. The model starts from a very good initial point (the pretrained weights), and the goal is to adapt the weights gently to the task without disrupting the broad representations learned during pretraining.

For fine-tuning, shorter warmup periods are generally appropriate. A common recommendation is 6-10% of total fine-tuning steps for warmup, which is higher as a fraction of total steps than in pretraining, because the total number of fine-tuning steps is much smaller (often 1,000-10,000 rather than 100,000-1,000,000). In absolute terms, fine-tuning warmup is often only a few hundred steps.

The peak learning rate for fine-tuning should also be much smaller than the pretraining peak: typically $10^{-5}$ to $5 \times 10^{-5}$ rather than $10^{-4}$ to $3 \times 10^{-4}$. This conservative learning rate reduces the risk of catastrophic forgetting (covered in the Catastrophic Forgetting chapter), where fine-tuning on the task overwrites the general representations acquired during pretraining.

Instruction tuning, which trains a model to follow natural language instructions by exposing it to instruction-following demonstrations, has become a standard step in the development of conversational AI systems. Instruction tuning is a form of supervised fine-tuning, but it typically uses a much more diverse set of tasks than traditional fine-tuning, including summarization, question answering, code generation, and mathematical reasoning.

For instruction tuning, the warmup configuration is influenced by the number of gradient steps and the peak learning rate. Instruction tuning runs are often very short (a few thousand to tens of thousands of steps), so warmup of 100-500 steps is typical. The learning rate is kept conservative to prevent the instruction tuning from washing out the pretraining representations. The combination of short warmup, low peak learning rate, and gradual decay has been validated across many instruction-tuning setups including the original InstructGPT experiments and subsequent work.