Learning Rate Schedulers in Nanotron: Warm-Up and Decay Configuration Guide

Nanotron implements a flexible LambdaLR-based learning rate scheduler that supports linear and constant warm-up styles combined with linear, cosine, or 1-sqrt decay curves, all configurable through the LRSchedulerArgs dataclass.

Nanotron provides a configurable learning rate scheduling system designed for large-scale language model training. The scheduler is built on PyTorch’s LambdaLR and allows precise control over warm-up and decay phases through the LRSchedulerArgs configuration class in src/nanotron/config/config.py. This implementation supports per-parameter-group scheduling and integrates directly into the training loop via the lr_scheduler_builder function.

Available Learning Rate Scheduler Components

Nanotron’s scheduler consists of three configurable phases: warm-up, decay, and a fixed-rate mode. These are defined in the LRSchedulerArgs dataclass and processed by the builder logic in src/nanotron/helpers.py.

Warm-Up Styles

The scheduler supports two warm-up strategies that determine how the learning rate reaches its target value:

• Linear (default): Ramps the learning rate from 0 to the target learning_rate over lr_warmup_steps steps. According to the source in src/nanotron/helpers.py (lines 22-27), the computation follows initial_lr * current_step / max(lr_warmup_steps, 1).

• Constant: Maintains the learning rate at the target value throughout the warm-up phase. This is implemented by returning lr_scheduler_args.learning_rate directly during the warm-up period.

Decay Styles

After warm-up completes, the scheduler transitions to one of three decay curves that reduce the learning rate toward min_decay_lr:

• Linear (default): Decays the learning rate linearly from the initial value to min_decay_lr over lr_decay_steps. The formula (lines 35-40 in helpers.py) calculates the remaining progress as (lr_decay_steps - (current_step - lr_decay_starting_step)) / lr_decay_steps.

• Cosine: Applies a cosine annealing curve (lines 41-44) following (1 + cos(π * progress)) / 2 to smoothly transition between the initial and minimum learning rates.

• 1-sqrt: Uses a square-root decay curve (lines 45-53) calculated as 1 - sqrt((current_step - lr_decay_starting_step) / lr_decay_steps), providing a sharper initial drop that gradually flattens.

Fixed Learning Rate Mode

When both lr_warmup_steps and lr_decay_steps are set to 0, the scheduler returns the initial learning rate unchanged for every step. This effectively disables scheduling and maintains a constant learning rate throughout training.

Implementation Architecture

The learning rate scheduler architecture in Nanotron separates configuration from execution. The LRSchedulerArgs dataclass in src/nanotron/config/config.py (lines 54-71) defines the parameters, while src/nanotron/helpers.py contains the lr_scheduler_builder function (lines 79-84) that constructs the actual PyTorch scheduler.

The builder creates a custom lr_lambda function for each optimizer parameter group. This lambda computes the multiplicative factor for the learning rate at each step, accounting for warm-up progress, decay progress, and normalization. The implementation normalizes the final lambda by dividing by the initial learning rate (line 61) to match PyTorch’s expected lr_lambda signature, which expects a multiplicative factor rather than an absolute value.

The resulting torch.optim.lr_scheduler.LambdaLR instance is attached to the trainer in src/nanotron/trainer.py at line 225, ensuring the scheduler updates automatically during the training loop.

Configuration Examples

Linear Warm-Up with Cosine Decay

This configuration ramps the learning rate linearly over the first 500 steps, then applies cosine decay from 1e-4 down to 1e-5 over the next 9,500 steps:

from nanotron.config.config import LRSchedulerArgs

lr_sched = LRSchedulerArgs(
    learning_rate=1e-4,
    lr_warmup_steps=500,
    lr_warmup_style="linear",
    lr_decay_style="cosine",
    lr_decay_steps=9500,
    min_decay_lr=1e-5,
)

Constant Warm-Up with 1-Sqrt Decay

To skip warm-up entirely and apply a square-root decay curve starting from step zero:

lr_sched = LRSchedulerArgs(
    learning_rate=2e-4,
    lr_warmup_steps=0,
    lr_warmup_style="constant",
    lr_decay_style="1-sqrt",
    lr_decay_steps=20000,
    min_decay_lr=5e-5,
)

Disabling the Scheduler for Fixed Learning Rate

To maintain a constant learning rate of 3e-4 throughout training without any warm-up or decay:

lr_sched = LRSchedulerArgs(
    learning_rate=3e-4,
    lr_warmup_steps=0,
    lr_decay_steps=0,
)

Integration with the Training Loop

The scheduler is instantiated within the Trainer class initialization in src/nanotron/trainer.py (line 225). The lr_scheduler_builder function receives the optimizer, scheduler arguments, and total training steps to construct the LambdaLR instance:

from nanotron.helpers import lr_scheduler_builder
from nanotron.trainer import Trainer

# Inside Trainer.__init__ (trainer.py line 225)

self.lr_scheduler = lr_scheduler_builder(
    optimizer=self.optimizer,
    lr_scheduler_args=self.config.optimizer.learning_rate_scheduler,
    total_training_steps=self.config.tokens.train_steps,
)

This integration ensures the learning rate updates automatically at each optimization step according to the configured warm-up and decay schedules.

Summary

• Nanotron’s scheduler is implemented as a LambdaLR-based system in src/nanotron/helpers.py with configuration defined in src/nanotron/config/config.py.
• Warm-up options include linear ramping and constant hold styles.
• Decay options support linear, cosine, and 1-sqrt curves that reduce learning rate toward a minimum threshold.
• The scheduler handles per-parameter-group learning rates and normalizes lambda values by the initial learning rate to match PyTorch conventions.
• Fixed learning rate training is achieved by setting both warm-up and decay steps to zero.

Frequently Asked Questions

What warm-up styles does Nanotron support?

Nanotron supports two warm-up styles: linear (default) and constant. Linear warm-up ramps the learning rate from zero to the target value over the specified steps, while constant warm-up maintains the target learning rate throughout the warm-up phase. These are configured via the lr_warmup_style parameter in LRSchedulerArgs.

How do I configure cosine decay in Nanotron?

Set lr_decay_style="cosine" in your LRSchedulerArgs configuration. The scheduler will then follow a cosine annealing curve from the initial learning rate down to min_decay_lr over the specified lr_decay_steps. The cosine calculation is implemented in src/nanotron/helpers.py using the formula (1 + cos(π * progress)) / 2.

Can I use a constant learning rate without warm-up or decay?

Yes. To disable scheduling and use a fixed learning rate, set both lr_warmup_steps=0 and lr_decay_steps=0 in your LRSchedulerArgs. When both values are zero, the lr_scheduler_builder returns a scheduler that maintains the initial learning rate unchanged for all training steps.

Where is the learning rate scheduler implemented in the Nanotron codebase?

The core logic resides in src/nanotron/helpers.py within the lr_scheduler_builder function (lines 79-84), which constructs the PyTorch LambdaLR scheduler. The configuration dataclass LRSchedulerArgs is defined in src/nanotron/config/config.py (lines 54-71). The scheduler is integrated into the training loop in src/nanotron/trainer.py at line 225.

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