# Implementing Learning Rate Schedulers for LLM Pretraining: Warm-Up and Cosine Decay in PyTorch

> Implement learning rate schedulers for LLM pretraining with PyTorch. Combine linear warm-up and cosine decay for stable convergence in transformer models. Learn how from LLMs-from-scratch.

- Repository: [Sebastian Raschka/LLMs-from-scratch](https://github.com/rasbt/LLMs-from-scratch)
- Tags: how-to-guide
- Published: 2026-05-12

---

**The rasbt/LLMs-from-scratch repository implements a two-phase learning rate scheduler that combines linear warm-up with cosine annealing directly inside the training loop, eliminating external dependencies while ensuring stable convergence for transformer models.**

Implementing learning rate schedulers for LLM pretraining is essential for preventing early-training instabilities and maximizing final model performance. Rather than relying on external scheduling libraries, the rasbt/LLMs-from-scratch repository embeds a lightweight warm-up and cosine decay mechanism directly into the training orchestration code, providing full transparency and control over the optimization trajectory.

## How the Learning Rate Scheduler Works

The scheduler operates in two distinct phases that mirror the optimization strategies used in modern large language models like GPT-3. Implemented in [`pkg/llms_from_scratch/appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/pkg/llms_from_scratch/appendix_d.py), the logic automatically adjusts the optimizer's learning rate at every global step based on the current training progress.

### Linear Warm-Up Phase

During the initial `warmup_steps`, the learning rate increases linearly from an `initial_lr` to the peak learning rate defined in the optimizer configuration. According to lines 36-38 in [`appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/appendix_d.py), the increment per step is calculated as:

```python
lr_increment = (peak_lr - initial_lr) / warmup_steps
current_lr = initial_lr + global_step * lr_increment

```

This gradual ramp prevents large gradient spikes during the early stages of training when the model weights are far from optimal. The peak learning rate is extracted dynamically from `optimizer.param_groups[0]["lr"]`, ensuring compatibility with any PyTorch optimizer that respects the standard parameter group interface.

### Cosine Annealing Phase

Once the warm-up phase completes (`global_step >= warmup_steps`), the scheduler transitions to a cosine decay curve that smoothly reduces the learning rate from the peak down to a specified `min_lr`. As implemented in lines 50-54 of [`appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/appendix_d.py), the calculation follows:

```python
progress = (global_step - warmup_steps) / (total_steps - warmup_steps)
current_lr = min_lr + (peak_lr - min_lr) * 0.5 * (1 + cos(pi * progress))

```

This cosine annealing approach often yields better final perplexity scores compared to step-wise or exponential decay schedules, as it allows for fine-grained weight updates during the later stages of training when the model is approaching convergence.

## Integration with the Training Loop

The `train_model` function in [`pkg/llms_from_scratch/appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/pkg/llms_from_scratch/appendix_d.py) serves as the central orchestrator, handling not only the forward and backward passes but also the learning rate updates and gradient management. The scheduler writes the computed learning rate back to every parameter group via `param_group["lr"] = current_lr`, making it transparent to the underlying optimizer.

Additionally, the implementation includes gradient clipping (lines 64-70) to prevent exploding gradients. Note that the corrected version uses `>=` rather than `>` for the clipping threshold comparison, ensuring proper handling of boundary cases during the warm-up phase.

## Implementation Examples

### Basic Usage with train_model

To utilize the scheduler in a standard training script such as [`ch05/01_main-chapter-code/gpt_train.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/01_main-chapter-code/gpt_train.py), pass the warm-up and learning rate bounds directly to the training function:

```python
import torch
from pkg.llms_from_scratch.appendix_d import train_model

# Initialize model and optimizer

model = MyTransformer(vocab_size=50257, d_model=768, n_heads=12, ...)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)  # This becomes peak_lr

# Configure scheduler parameters

warmup_steps = 500
initial_lr = 3e-5
min_lr = 1e-6

# Execute training with automatic LR scheduling

train_losses, val_losses, tokens_seen, track_lrs = train_model(
    model=model,
    train_loader=train_loader,
    val_loader=val_loader,
    optimizer=optimizer,
    device="cuda",
    n_epochs=10,
    eval_freq=1000,
    warmup_steps=warmup_steps,
    initial_lr=initial_lr,
    min_lr=min_lr,
    orig_book_version=False,  # Use corrected gradient clipping

)

```

The function returns `track_lrs`, a list containing the learning rate value at every optimization step, which enables precise monitoring and debugging of the schedule behavior.

### Hyperparameter Search Over Warm-Up Steps

For systematic exploration of optimal warm-up durations, the repository provides [`ch05/05_bonus_hparam_tuning/hparam_search.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/05_bonus_hparam_tuning/hparam_search.py). This utility accepts a grid configuration to test different `warmup_iters` values:

```python
from ch05_05_bonus_hparam_tuning.hparam_search import run_hparam_search

hparam_grid = {
    "warmup_iters": [100, 250, 500, 1000],
    "learning_rate": [1e-4, 5e-4, 1e-3],
}

run_hparam_search(
    train_loader=train_loader,
    val_loader=val_loader,
    model_factory=lambda: MyTransformerModel(),
    optimizer_factory=lambda model: torch.optim.AdamW(model.parameters(), lr=5e-4),
    hparam_config=hparam_grid,
    max_epochs=3,
)

```

Each trial automatically instantiates the scheduler with the specified `warmup_iters`, allowing you to identify the configuration that minimizes validation loss for your specific dataset and model architecture.

### Monitoring the Learning Rate Schedule

Visualizing the learning rate progression helps verify that the warm-up and decay phases are operating as expected:

```python
import matplotlib.pyplot as plt

# track_lrs is returned by train_model

plt.plot(track_lrs)
plt.title("Learning Rate Schedule: Warm-Up + Cosine Decay")
plt.xlabel("Training Step")
plt.ylabel("Learning Rate")
plt.axvline(x=warmup_steps, color='r', linestyle='--', label='End of Warm-Up')
plt.legend()
plt.show()

```

This inspection is particularly valuable when scaling to distributed training scenarios, such as those implemented in [`ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py), where ensuring consistent learning rate updates across processes is critical.

## Key Files and Source Locations

- **[`pkg/llms_from_scratch/appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/pkg/llms_from_scratch/appendix_d.py)** (lines 36-70): Core implementation of the `train_model` function containing the warm-up calculation (lines 36-38), cosine annealing logic (lines 50-54), and gradient clipping (lines 64-70).

- **[`ch05/01_main-chapter-code/gpt_train.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/01_main-chapter-code/gpt_train.py)**: Reference single-GPU training script demonstrating how to invoke `train_model` with scheduler arguments.

- **[`ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py)**: Distributed Data Parallel (DDP) implementation showcasing the scheduler in a multi-GPU environment.

- **[`ch05/05_bonus_hparam_tuning/hparam_search.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/05_bonus_hparam_tuning/hparam_search.py)**: Hyperparameter search utility for systematically evaluating different `warmup_iters` configurations.

- **[`pkg/llms_from_scratch/tests/test_appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/pkg/llms_from_scratch/tests/test_appendix_d.py)**: Unit tests verifying the mathematical correctness of the warm-up and cosine decay calculations.

## Summary

- The rasbt/LLMs-from-scratch repository implements a **linear warm-up with cosine annealing** schedule directly in [`pkg/llms_from_scratch/appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/pkg/llms_from_scratch/appendix_d.py), requiring no external scheduler classes.
- The scheduler automatically extracts the **peak learning rate** from the optimizer's initial configuration and interpolates between `initial_lr`, `peak_lr`, and `min_lr` based on `global_step`.
- **Gradient clipping** is integrated into the same training loop (lines 64-70), with the corrected version using proper boundary checking to ensure numerical stability.
- The `train_model` function returns a `track_lrs` list that enables precise monitoring and visualization of the learning rate trajectory.
- **Hyperparameter search** capabilities in [`hparam_search.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/hparam_search.py) allow systematic exploration of warm-up durations to optimize training dynamics for specific hardware and datasets.

## Frequently Asked Questions

### Why combine warm-up with cosine annealing for LLM pretraining?

**Warm-up stabilizes early training** by preventing large gradient steps when the model is initialized with random weights, while **cosine annealing** provides a smooth decay that allows the optimizer to settle into flat minima rather than bouncing between sharp ones. This combination is empirically shown to produce lower final perplexity than constant or step-decay schedules in transformer architectures.

### How do I adjust the warm-up duration for different model sizes?

Modify the `warmup_steps` parameter passed to `train_model` or the `warmup_iters` key in the hyperparameter search grid. As a rule of thumb documented in the repository's training scripts, larger models typically benefit from longer warm-up periods (500-2000 steps) to accommodate more complex loss landscapes, while smaller models may converge with as few as 100-250 warm-up steps.

### Can I use this scheduler with optimizers other than AdamW?

Yes. The scheduler is **optimizer-agnostic** because it writes the computed learning rate directly to `param_group["lr"]` for every parameter group in the optimizer. This approach works with any PyTorch optimizer that respects the learning rate field in its parameter groups, including SGD, Adam, and AdamW variants.

### How does gradient clipping interact with the learning rate schedule?

Gradient clipping (applied at lines 64-70 in [`appendix_d.py`](https://github.com/rasbt/LLMs-from-scratch/blob/main/appendix_d.py)) operates **after** the learning rate update, ensuring that the scaled gradients remain within a norm threshold regardless of the current learning rate magnitude. This sequencing prevents the large parameter updates that can occur during the transition from warm-up to cosine phases, maintaining training stability throughout the optimization process.