# How to Configure Learning Rate Decay for Transformer Training in train-llm-from-scratch

> Learn to configure learning rate decay for Transformer training in train-llm-from-scratch by setting T_LR_DECAY_STEP and T_LR_DECAYED in config python. Optimize your model training.

- Repository: [Fareed Khan/train-llm-from-scratch](https://github.com/FareedKhan-dev/train-llm-from-scratch)
- Tags: how-to-guide
- Published: 2026-05-31

---

**You configure learning rate decay by setting `T_LR_DECAY_STEP` and `T_LR_DECAYED` in [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py), which the training loop in [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py) applies as a single-step drop at the specified step.**

The `train-llm-from-scratch` repository implements a straightforward single-step learning rate decay mechanism designed for transformer language model training. Unlike complex scheduling libraries, this implementation uses a simple equality check inside the training loop to drop the learning rate at a specific step. Understanding how to configure learning rate decay for transformer training in this codebase requires modifying central configuration constants and understanding how the optimizer interacts with the training script.

## Where Learning Rate Decay is Configured

The repository separates hyperparameter definitions from training logic, making adjustments possible without touching the core training algorithm.

### Configuration Constants in config/config.py

All decay-related settings live in [[`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py)](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py). The file exports three critical constants through the `default_config` dictionary:

- **`T_LR_DECAY_STEP`**: The exact training step (integer) when decay occurs
- **`T_LR`**: The initial learning rate (e.g., `5e-4`)
- **`T_LR_DECAYED`**: The target learning rate after decay (e.g., `5e-5`)

These map to dictionary keys `'t_lr_decay_step'`, `'t_lr'`, and `'t_lr_decayed'` respectively, which the training script consumes.

### Optimizer Setup in scripts/train_transformer.py

The [[`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py)](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py) file instantiates an **AdamW** optimizer using the initial rate from `config['t_lr']`. The optimizer object maintains the learning rate state across all parameter groups, which the decay logic modifies directly during training.

## How the Decay Mechanism Works

The implementation uses a hard step equality rather than a gradual schedule. Inside the main training loop, the code checks `if step == config['t_lr_decay_step']:` at every iteration. When this condition triggers, the script prints "Decaying learning rate" and overwrites the `lr` field for **all** parameter groups in the optimizer with `config['t_lr_decayed']`.

This design means:
- The decay fires **exactly once** during the entire training run
- The drop is immediate, not gradual
- All optimizer parameter groups receive the same new learning rate

## Practical Configuration Examples

### Default Single-Step Decay

To use the repository's built-in decay behavior, edit [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py) with your desired timing:

```python

# config/config.py

T_LR_DECAY_STEP = 50000      # Decay after 50k steps

T_LR = 5e-4                  # Start at 0.0005

T_LR_DECAYED = 5e-5          # Drop to 0.00005

```

Run the trainer normally; the optimizer automatically switches to the decayed rate at step 50,000 without requiring code changes in the training script.

### Customizing Decay Timing and Values

Adjust the constants to match your convergence behavior:

```python

# config/config.py

T_LR_DECAY_STEP = 100000     # Delay decay until 100k steps

T_LR = 1e-3                  # Higher initial rate

T_LR_DECAYED = 1e-4          # Gentler drop to 0.0001

```

Because [`train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/train_transformer.py) reads these values from the config dictionary at runtime, you can tune decay behavior purely through configuration changes.

### Implementing Advanced Schedules

If you require smooth decay (linear, cosine, or exponential) rather than a single step, modify the decay logic in [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py). Replace the equality check with interpolation logic:

```python

# Inside the training loop, after optimizer.step()

if step < config['t_lr_decay_step']:
    # Linear warm-up phase

    warmup_factor = step / config['t_lr_decay_step']
    lr = config['t_lr'] * warmup_factor
else:
    # Linear decay phase to final value

    progress = (step - config['t_lr_decay_step']) / (config['t_train_steps'] - config['t_lr_decay_step'])
    lr = config['t_lr'] * (1 - progress) + config['t_lr_decayed'] * progress

for g in optimizer.param_groups:
    g['lr'] = lr

```

This modification transitions smoothly from the initial rate to the decayed rate over the remaining training steps.

## Summary

- **Configuration location**: Set `T_LR_DECAY_STEP` and `T_LR_DECAYED` in [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py) to control when and to what value the learning rate drops
- **Mechanism**: [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py) triggers decay via an equality check (`step == config['t_lr_decay_step']`) that fires once per training run
- **Default behavior**: Immediate single-step reduction applied to all AdamW parameter groups simultaneously
- **Extensibility**: Replace the equality check with custom interpolation logic for cosine, exponential, or linear schedules

## Frequently Asked Questions

### What is the default decay behavior?

The default configuration performs a **single-step decay** where the learning rate drops instantly from `T_LR` to `T_LR_DECAYED` at the exact step specified by `T_LR_DECAY_STEP`. This occurs once during training and affects all parameter groups in the AdamW optimizer equally.

### Can I implement cosine annealing or exponential decay?

Yes, but you must modify the training loop logic in [`scripts/train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py). The current implementation only supports single-step drops via an equality check. To implement cosine annealing, replace the `if step == config['t_lr_decay_step']:` block with a function that calculates the current learning rate based on training progress and updates `optimizer.param_groups` accordingly at every step.

### Why does the decay use an equality check instead of a conditional?

The equality check (`step == config['t_lr_decay_step']`) ensures the decay happens **exactly once** at a predetermined point. This design choice supports training pipelines where you want to manually specify convergence milestones rather than following a mathematical curve. It also prevents unnecessary computation during steps where no decay occurs.

### Do I need to restart training to change the decay step?

No, you only need to edit [`config/config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/config.py) and relaunch the training script. Because [`train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/train_transformer.py) reads `t_lr_decay_step` from the config dictionary at initialization, changes to `T_LR_DECAY_STEP` take effect immediately on the next training run. If resuming from a checkpoint, ensure the saved step count hasn't already passed your new decay step value.