# How to Resume Training from a Checkpoint in Nanotron

> Easily resume training from a checkpoint in Nanotron. Set resume_checkpoint_path and run train.py to seamlessly continue your model training.

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

---

**To resume training in Nanotron, set the `resume_checkpoint_path` parameter in your `CheckpointsArgs` configuration and launch training with [`run_train.py`](https://github.com/huggingface/nanotron/blob/main/run_train.py); the framework automatically loads model weights, optimizer state, and determines the correct data stage to continue from.**

Resuming long-running distributed training jobs is essential for fault tolerance and iterative experimentation. In the `huggingface/nanotron` repository, checkpoint resumption is implemented through a coordinated pipeline involving configuration parsing, filesystem discovery, and data stage reconciliation.

## Checkpoint Resumption Architecture

Nanotron implements resume training through three tightly coupled components that handle state restoration before the training loop restarts.

### Configuration via CheckpointsArgs

The resume process begins in [`src/nanotron/config/config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py), where the `CheckpointsArgs` dataclass defines the `resume_checkpoint_path` parameter. When this field points to a valid checkpoint directory, Nanotron automatically triggers the loading sequence for model parameters, optimizer buffers, and learning rate scheduler state.

You can control exactly what gets loaded by toggling boolean flags in the same configuration object:

- `load_optimizer` – Restore optimizer states (default: `True`)
- `load_lr_scheduler` – Restore learning rate schedule position (default: `True`)

### Checkpoint Discovery with parse_ckpt_path

Once the configuration is parsed, `nanotron.serialize.main.parse_ckpt_path` (located in [`src/nanotron/serialize/main.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/serialize/main.py)) handles the actual filesystem operations. This function examines the supplied `resume_checkpoint_path` and resolves the concrete checkpoint location through the following logic:

- **Local checkpoints**: Searches for [`latest.txt`](https://github.com/huggingface/nanotron/blob/main/latest.txt) (containing the iteration number) or [`model_config.json`](https://github.com/huggingface/nanotron/blob/main/model_config.json) to identify the most recent save
- **S3 checkpoints**: Downloads the checkpoint referenced by [`latest.txt`](https://github.com/huggingface/nanotron/blob/main/latest.txt) if the path is a remote URI

The function returns the resolved filesystem path that the serializer uses to load tensors and metadata.

### Data Stage Resumption Logic

After model weights are restored, the trainer must determine which data stage to continue from. The `Trainer.find_stage_idx_to_resume` method in [`src/nanotron/trainer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/trainer.py) scans `config.data_stages` in reverse order to locate the most recent stage whose `start_training_step` is less than or equal to the current global step. 

The trainer then treats that stage as a "resume" stage, recalculates remaining training steps, and logs consumption statistics (samples processed versus remaining).

## Step-by-Step Resume Workflow

When you launch [`run_train.py`](https://github.com/huggingface/nanotron/blob/main/run_train.py) with a resume configuration, Nanotron executes the following sequence:

1. **Configuration Loading**: The YAML config containing `resume_checkpoint_path` is parsed into a `Config` object with populated `CheckpointsArgs`
2. **Checkpoint Resolution**: `parse_ckpt_path` locates the latest checkpoint files (weights, optimizer state, metadata) locally or downloads from S3
3. **State Restoration**: The serializer loads model weights, and optionally optimizer and LR scheduler states based on your `CheckpointsArgs` flags
4. **Stage Calculation**: `find_stage_idx_to_resume` identifies the correct data stage using the restored global step count and `last_stage_idx` from checkpoint metadata
5. **Training Continuation**: The trainer reinitializes data loaders at the correct sample offset and resumes the training loop

## Configuration Example

To resume training from an existing checkpoint, instantiate your `Config` with `CheckpointsArgs` pointing to the previous run directory:

```python
from nanotron.config import (
    Config, CheckpointsArgs, GeneralArgs, ParallelismArgs, 
    ModelArgs, OptimizerArgs, LRSchedulerArgs, DataArgs
)

CHECKPOINT_PATH = "./checkpoints/smollm2-135m-nanotron"

config = Config(
    general=GeneralArgs(project="resume_demo", run="resume_%date_%jobid", seed=42),
    checkpoints=CheckpointsArgs(
        checkpoints_path="./checkpoints",
        checkpoint_interval=10,
        resume_checkpoint_path=CHECKPOINT_PATH,  # Path to previous checkpoint

        load_optimizer=True,                      # Load optimizer state

        load_lr_scheduler=True,                   # Load LR scheduler state

    ),
    parallelism=ParallelismArgs(dp=2, pp=1, tp=1),
    model=ModelArgs(init_method=your_init, model_config=your_model_cfg),
    optimizer=OptimizerArgs(
        zero_stage=0,
        learning_rate_scheduler=LRSchedulerArgs(
            learning_rate=3e-4, lr_warmup_steps=2, 
            lr_warmup_style="linear", lr_decay_style="cosine"
        ),
    ),
    data_stages=[your_stage_config],
)

config.save_as_yaml("resume_config.yaml")

```

This pattern mirrors the official example in [`examples/config_resume_training.py`](https://github.com/huggingface/nanotron/blob/main/examples/config_resume_training.py).

## Launching Resumed Training

Execute the training script with your resume configuration file:

```bash
python run_train.py --config resume_config.yaml

```

The `NanotronTrainer` automatically detects the `resume_checkpoint_path` setting and orchestrates the full restoration pipeline before processing the first training batch.

## Important Constraints and Considerations

Resuming from checkpoints in Nanotron imposes specific technical constraints that prevent configuration drift.

**Same Parallelism Required**: The checkpoint can only be resumed when the tensor-parallel size (`tp`) matches the size used when the checkpoint was created. This validation is enforced during optimizer state deserialization in [`src/nanotron/serialize/optimizer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/serialize/optimizer.py).

**Selective Loading**: Set `load_optimizer=False` or `load_lr_scheduler=False` in `CheckpointsArgs` if you only want to restore model weights while resetting the optimizer or learning rate schedule.

**Metadata Consistency**: The checkpoint metadata stores `last_stage_idx` and `last_train_step`. Changing the data-stage schedule (e.g., inserting new stages or modifying `start_training_step` values) after creating the checkpoint may break the resume logic and cause the trainer to select an incorrect stage.

## Summary

- Set `resume_checkpoint_path` in `CheckpointsArgs` to enable automatic checkpoint loading
- `parse_ckpt_path` in [`src/nanotron/serialize/main.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/serialize/main.py) handles local and S3 checkpoint discovery
- The trainer uses `find_stage_idx_to_resume` to synchronize data stages with the restored global step
- Tensor parallel dimensions must match between the checkpoint and current configuration
- Use `load_optimizer` and `load_lr_scheduler` flags to control selective state restoration

## Frequently Asked Questions

### Can I resume training with different tensor parallel settings?

No. Nanotron enforces that the tensor-parallel size (`tp`) used when creating the checkpoint must exactly match the `tp` setting in your resume configuration. Attempting to load with mismatched parallelism dimensions will raise an error during optimizer state restoration in [`src/nanotron/serialize/optimizer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/serialize/optimizer.py).

### How do I resume training without loading the optimizer state?

Set `load_optimizer=False` in your `CheckpointsArgs` configuration. This restores only the model weights while initializing a fresh optimizer state, which is useful when changing optimization hyperparameters or continuing training with a different learning rate schedule.

### What happens if I modify my data stages after creating a checkpoint?

Modifying the `data_stages` configuration (such as inserting new stages or changing `start_training_step` values) may break the resume logic. The trainer relies on `last_stage_idx` metadata stored in the checkpoint to locate the correct stage via `find_stage_idx_to_resume`. If the stage indices no longer align with the restored global step, the trainer may select an incorrect stage or fail to resume properly.

### Does Nanotron support resuming from S3 checkpoints?

Yes. When `resume_checkpoint_path` points to an S3 URI, `parse_ckpt_path` automatically downloads the checkpoint files referenced by [`latest.txt`](https://github.com/huggingface/nanotron/blob/main/latest.txt) before loading them into the model and optimizer. Ensure your environment has appropriate AWS credentials configured for the download operation.