# How to Set Up Multi-Stage Dataset Training in Nanotron

> Learn how to set up multi-stage dataset training in Nanotron. Configure sequential training phases and automatic dataloader switching with the data_stages field for efficient model development.

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

---

**Nanotron enables multi-stage dataset training through the `data_stages` configuration field, which allows you to define sequential training phases with automatic dataloader switching at specified global steps.**

The HuggingFace Nanotron framework supports complex training curricula where you transition from pre-training data to instruction-tuning or domain-specific datasets without restarting the training run. By configuring the `data_stages` attribute in your `Config` object, you can specify exactly when each dataset becomes active, optionally override sequence lengths per stage, and resume training seamlessly from checkpoints.

## Defining Stages in the Configuration Schema

Multi-stage training is controlled by the `data_stages` field in [`src/nanotron/config/config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py). Each stage is defined as a `DatasetStageArgs` instance that requires four key parameters:

- **`name`** – A unique string identifier for the stage (e.g., `"pretrain"` or `"sft"`).
- **`start_training_step`** – The global training step at which the stage begins. The first stage must start at step `1`, and subsequent stages must have unique, increasing step values.
- **`data`** – A `DataArgs` object pointing to your dataset configuration (e.g., `PretrainDatasetsArgs`, `NanosetDatasetsArgs`, or `SFTDatasetsArgs`).
- **`sequence_length`** – Optional per-stage sequence length that overrides the global `tokens.sequence_length`.

During initialization, the `Config` class validates these stages in `__post_init__` (lines 90-106) to ensure unique names and start steps, then sorts them by `start_training_step` to establish execution order. This validation prevents configuration errors before training begins.

## Stage Transition Mechanics

The `NanotronTrainer` class in [`src/nanotron/trainer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/trainer.py) handles stage transitions automatically through the `_update_dataloader_based_on_training_stages` method (lines 24-84). This mechanism operates every training iteration to determine which dataset should be active.

### Stage Selection Logic

The trainer identifies the active stage by scanning for the first stage where `start_training_step` equals the current `iteration_step`. When resuming from a checkpoint, it calculates which stage should continue based on previously completed steps. The trainer maintains the current stage state in `self.current_dataloader` and `self.current_base_dl` after performing sanity checks on the dataloader mapping.

### Memory Management Between Stages

When a stage boundary is reached, the trainer explicitly clears the previous stage's resources to prevent out-of-memory errors. The implementation deletes the previous dataloader object and its underlying datasets, then forces Python garbage collection (lines 46-62). This cleanup is critical when transitioning between large datasets like The Pile and smaller instruction-tuning sets, as it frees GPU memory that would otherwise remain allocated.

## Resuming Multi-Stage Training

Nanotron provides dedicated helper functions in [`src/nanotron/helpers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/helpers.py) (lines 797-823) to handle checkpoint resumption correctly:

- **`compute_remain_train_steps_of_a_data_stage_from_ckp`** – Calculates how many training steps remain in the current stage based on the checkpoint's global step counter.
- **`get_consumed_train_samples_of_a_data_stage_from_ckp`** – Retrieves the number of samples already processed in the current stage to maintain accurate data sampling.

These utilities update the `TrainingMetadata` object so that training continues exactly where it left off, preserving random states and data loader positions across stage boundaries.

## Implementation Examples

### YAML Configuration

Define your stages in the configuration file using the `data_stages` list:

```yaml
general:
  project: "multistage-llm"
  run: "run1"

tokens:
  micro_batch_size: 4
  batch_accumulation_per_replica: 2
  sequence_length: 2048
  train_steps: 150_000

data_stages:
  - name: "pretrain"
    start_training_step: 1
    data:
      dataset:
        hf_dataset_or_datasets: "EleutherAI/pile"
        hf_dataset_splits: "train"
        text_column_name: "text"
  
  - name: "sft"
    start_training_step: 100_000
    data:
      dataset:
        hf_dataset_or_datasets: "my-team/instruction-data"
        hf_dataset_splits: ["train", "validation"]
        text_column_name: "instruction"

```

Load this configuration using `Config.from_yaml()` or the CLI launcher. The framework automatically handles the transition at step 100,000.

### Python Training Setup

Instantiate the trainer with a mapping of stage names to dataloaders:

```python
from nanotron.config.config import Config, get_config_from_yaml
from nanotron.trainer import NanotronTrainer
from nanotron.utils import init_distributed

# Initialize distributed environment

world_info = init_distributed()

# Load configuration

cfg: Config = get_config_from_yaml("multistage_config.yaml")

# Build model and optimizer (using Nanotron utilities)

model = build_model(cfg)
optimizer = build_optimizer(cfg, model.parameters())

# Create stage-to-dataloader mapping

dataloaders = {
    "pretrain": lambda: make_dataloader(cfg.data_stages[0].data.dataset, cfg),
    "sft": lambda: make_dataloader(cfg.data_stages[1].data.dataset, cfg),
}

trainer = NanotronTrainer(
    config=cfg,
    model=model,
    optimizer=optimizer,
    parallel_context=world_info.parallel_context,
)

trainer.train(dataloader_or_dls=dataloaders)

```

The trainer expects either a single dataloader or a dictionary mapping stage names to callable functions that return `torch.utils.data.DataLoader` instances.

### Custom Nanoset Dataloader

For binary-encoded datasets using Nanotron's custom format:

```python
def make_nanoset_dataloader(stage_data, cfg):
    from nanotron.dataset.nanoset import NanosetDataset, NanosetCollator
    
    dataset = NanosetDataset(
        folder=stage_data.dataset_folder,
        tokenizer_name=stage_data.tokenizer_name,
        vocab_size=stage_data.vocab_size,
    )
    
    collator = NanosetCollator(
        pad_id=dataset.pad_token_id,
        max_seq_len=stage_data.sequence_length or cfg.tokens.sequence_length,
    )
    
    return torch.utils.data.DataLoader(
        dataset,
        batch_size=cfg.tokens.micro_batch_size,
        num_workers=stage_data.num_loading_workers,
        collate_fn=collator,
        shuffle=True,
    )

```

Pass this callable in your `dataloaders` dictionary to enable efficient binary data loading for specific stages.

## Summary

- **Multi-stage dataset training** in Nanotron uses the `data_stages` configuration field to define sequential training phases with unique start steps.
- The `DatasetStageArgs` class in [`src/nanotron/config/config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py) specifies stage names, start steps, dataset configurations, and optional sequence lengths.
- `NanotronTrainer._update_dataloader_based_on_training_stages` automatically switches dataloaders at stage boundaries and clears previous stage memory to prevent OOM errors.
- Helper functions in [`src/nanotron/helpers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/helpers.py) enable accurate checkpoint resumption by calculating remaining steps and consumed samples per stage.
- You provide a mapping of stage names to dataloader callables, allowing flexible integration of HuggingFace datasets, custom binary formats, or streaming data sources.

## Frequently Asked Questions

### What happens if two stages have the same start_training_step?

The `Config` class validation in [`src/nanotron/config/config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/config.py) (lines 94-105) raises an error during initialization. Each stage must have a unique `start_training_step` value to prevent ambiguous stage boundaries.

### Can I use different sequence lengths for different training stages?

Yes. Each `DatasetStageArgs` accepts an optional `sequence_length` parameter that overrides the global `tokens.sequence_length` for that specific stage. This is useful when transitioning from long-context pre-training to shorter-context fine-tuning.

### How does Nanotron handle checkpoint resumption in the middle of a stage?

The trainer uses `compute_remain_train_steps_of_a_data_stage_from_ckp` and `get_consumed_train_samples_of_a_data_stage_from_ckp` from [`src/nanotron/helpers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/helpers.py) to recalculate the remaining steps and consumed samples. It then updates the `TrainingMetadata` to resume data loading from the correct position without reprocessing already-seen samples.

### Is it possible to mix different dataset types across stages?

Absolutely. You can combine `PretrainDatasetsArgs` for early stages, `NanosetDatasetsArgs` for mid-training efficiency, and `SFTDatasetsArgs` for final fine-tuning in the same configuration. Each stage's `data` field accepts any valid `DataArgs` subclass, allowing heterogeneous data sources throughout the training run.