How to Resume Training from a Checkpoint in Nanotron

To resume training in Nanotron, set the resume_checkpoint_path parameter in your CheckpointsArgs configuration and launch training with 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, 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) 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 (containing the iteration number) or model_config.json to identify the most recent save
• S3 checkpoints: Downloads the checkpoint referenced by 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 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 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:

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.

Launching Resumed Training

Execute the training script with your resume configuration file:

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.

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 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.

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 before loading them into the model and optimizer. Ensure your environment has appropriate AWS credentials configured for the download operation.

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