Implementation:Microsoft DeepSpeedExamples Token Based LR Annealing

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Deep Learning, Learning Rate Scheduling
Last Updated	2026-02-07 12:00 GMT

Overview

Implements a token-based learning rate annealing scheduler that supports linear warmup followed by cosine, linear, or constant decay styles.

Description

The AnnealingLR class provides a learning rate scheduler designed for large-scale language model training where decay is governed by the number of consumed tokens rather than training steps alone. During the initial warmup phase, the learning rate increases linearly from zero to max_lr over a specified number of warmup steps. After warmup, the scheduler tracks a warmup_tokens count that marks the token consumption at the end of warmup, and subsequent decay is computed based on the ratio of consumed tokens (minus warmup tokens) to the total decay token budget.

Three decay styles are supported: constant (no decay after warmup), linear (linear interpolation from max_lr to min_lr), and cosine (half-cosine annealing from max_lr to min_lr). Once the consumed token count exceeds the decay token budget, the learning rate is clamped to min_lr.

The class includes full checkpoint support via state_dict() and load_state_dict() methods, with configurable behavior for overriding or restoring scheduler state from checkpoints. The load_state_dict method handles backward compatibility with older checkpoint formats that use legacy key names such as start_lr, warmup_iter, end_iter, and num_iters.

Usage

Use this scheduler for GPT fine-tuning or pretraining tasks where the learning rate decay should be tied to token consumption rather than fixed step counts. This is particularly useful when batch sizes or sequence lengths vary during training, as token-based scheduling provides more consistent training dynamics.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/data_efficiency/gpt_finetuning/learning_rates.py
• Lines: 1-169

Signature

class AnnealingLR(object):
    def __init__(self, optimizer, max_lr, min_lr,
                 warmup_steps, decay_tokens, decay_style,
                 use_checkpoint_lr_scheduler=True,
                 override_lr_scheduler=False):
        ...

    def get_lr(self):
        ...

    def step(self, increment, consumed_tokens):
        ...

    def state_dict(self):
        ...

    def load_state_dict(self, sd):
        ...

Import

from learning_rates import AnnealingLR

I/O Contract

Inputs

Name	Type	Required	Description
optimizer	torch.optim.Optimizer	Yes	The optimizer whose learning rate will be adjusted
max_lr	float	Yes	Maximum (peak) learning rate after warmup
min_lr	float	Yes	Minimum learning rate floor (must be >= 0)
warmup_steps	int	Yes	Number of steps for linear warmup phase
decay_tokens	int	Yes	Total number of tokens over which to decay the learning rate
decay_style	str	Yes	Decay style: 'constant', 'linear', or 'cosine'
use_checkpoint_lr_scheduler	bool	No	Whether to use the checkpoint's scheduler values (default: True)
override_lr_scheduler	bool	No	Whether to override checkpoint values with class values (default: False)

Outputs

Name	Type	Description
cur_lr	float	The current learning rate after calling step()
state_dict	dict	Serializable state dictionary for checkpointing

Usage Examples

from learning_rates import AnnealingLR

# Create the scheduler
scheduler = AnnealingLR(
    optimizer=optimizer,
    max_lr=1e-4,
    min_lr=1e-6,
    warmup_steps=1000,
    decay_tokens=1_000_000_000,
    decay_style='cosine'
)

# During training loop
for step, batch in enumerate(dataloader):
    loss = model(batch)
    loss.backward()
    optimizer.step()
    tokens_consumed = step * batch_size * seq_length
    scheduler.step(increment=1, consumed_tokens=tokens_consumed)
    print(f"Current LR: {scheduler.cur_lr}")