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

You configure learning rate decay by setting T_LR_DECAY_STEP and T_LR_DECAYED in config/config.py, which the training loop in 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). 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) 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 with your desired timing:

# 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:

# 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 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. Replace the equality check with interpolation logic:

# 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 to control when and to what value the learning rate drops
• Mechanism: 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. 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 and relaunch the training script. Because 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}