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

> Explore Nanotron's flexible learning rate schedulers. Configure linear/cosine decay with warm-up styles using LRSchedulerArgs for optimal training. Learn more now.

- Repository: [Hugging Face/nanotron](https://github.com/huggingface/nanotron)
- Tags: how-to-guide
- Published: 2026-03-03

---

**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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py) (lines 54-71) defines the parameters, while [`src/nanotron/helpers.py`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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:

```python
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:

```python
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:

```python
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`](https://github.com/huggingface/nanotron/blob/main/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:

```python
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`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/helpers.py) with configuration defined in [`src/nanotron/config/config.py`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/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`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py) (lines 54-71). The scheduler is integrated into the training loop in [`src/nanotron/trainer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/trainer.py) at line 225.