# How to Use a Custom DataLoader with Nanotron: Implementation Guide

> Learn how to use a custom dataloader with Nanotron. Implement a collator for specific batch dictionary formats and TensorPointers for efficient distributed training.

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

---

**Yes, you can use a custom DataLoader with Nanotron by implementing a collator that returns the specific batch dictionary format required by the distributed training pipeline, including TensorPointers for non-owner ranks.**

Nanotron is a lightweight distributed training framework for large language models developed by Hugging Face. While it provides default data loading utilities through `get_train_dataloader()`, the architecture explicitly supports custom DataLoader implementations provided they adhere to the strict batch format contract required by pipeline parallelism, tensor parallelism, and context parallelism.

## Understanding Nanotron's DataLoader Contract

Nanotron's training loop expects batches as dictionaries with specific keys and value types. According to the source code in [`src/nanotron/data/clm_collator.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/clm_collator.py), each batch must contain:

- **`input_ids`**: `torch.Tensor` or `TensorPointer` containing token IDs for model input
- **`input_mask`**: `torch.Tensor` or `TensorPointer` for the attention mask
- **`label_ids`**: `torch.Tensor` or `TensorPointer` for prediction targets
- **`label_mask`**: `torch.Tensor` or `TensorPointer` indicating loss-contributing positions
- **`position_ids`** (optional): Required when `use_position_ids=True` for packed sequence handling

The critical requirement involves **TensorPointer** objects from `nanotron.parallel.pipeline_parallel.tensor_pointer`. These placeholder objects indicate which pipeline-parallel rank owns the actual tensor data. Your custom collator must return TensorPointers on ranks that do not own the data, while returning actual tensors on owner ranks. Failure to provide these pointers causes the trainer to reference non-existent tensors and raise runtime errors.

## Implementing a Custom Collator

### The Required Batch Schema

Your collator must respect the micro-batch size (`micro_batch_size`) and sequence length (`sequence_length`) defined in your training configuration. The collator runs in [`src/nanotron/data/dataloader.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/dataloader.py) within the DataLoader's worker processes, but must be aware of the distributed topology through the `parallel_context`.

### Creating the Custom Collator Class

Subclass or replace the existing `DataCollatorForCLM` logic found in [`src/nanotron/data/clm_collator.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/clm_collator.py). The following implementation demonstrates the contract:

```python
import numpy as np
from nanotron.parallel.pipeline_parallel.tensor_pointer import TensorPointer
import nanotron.distributed as dist

class MyCustomCollator:
    """Collate examples into the dict format required by Nanotron."""
    def __init__(self, seq_len: int, input_pp_rank: int, output_pp_rank: int, parallel_context):
        self.seq_len = seq_len
        self.input_pp_rank = input_pp_rank
        self.output_pp_rank = output_pp_rank
        self.parallel_context = parallel_context

    def __call__(self, examples):
        # Stack input_ids from dataset examples (assumed numpy arrays)

        batch = np.vstack([ex["input_ids"] for ex in examples])  # (B, L+1)

        batch_size, _ = batch.shape
        
        # Default: TensorPointers for ranks that don't own the data

        out = {
            "input_ids": TensorPointer(group_rank=self.input_pp_rank),
            "input_mask": TensorPointer(group_rank=self.input_pp_rank),
            "label_ids": TensorPointer(group_rank=self.output_pp_rank),
            "label_mask": TensorPointer(group_rank=self.output_pp_rank),
        }
        
        cur_pp = dist.get_rank(self.parallel_context.pp_pg)
        
        # Only input rank provides input tensors

        if cur_pp == self.input_pp_rank:
            out["input_ids"] = batch[:, :-1]  # Drop last token

            out["input_mask"] = np.ones((batch_size, self.seq_len), dtype=np.bool_)
            
        # Only output rank provides label tensors

        if cur_pp == self.output_pp_rank:
            out["label_ids"] = batch[:, 1:]   # Shift left for next-token prediction

            out["label_mask"] = np.ones((batch_size, self.seq_len), dtype=np.bool_)
            
        return out

```

This mirrors the logic in `DataCollatorForCLM` but allows custom preprocessing logic while maintaining the required output schema.

## Building and Configuring the DataLoader

### Using the Distributed Sampler

Nanotron requires proper dataset sharding across data-parallel ranks. Use `get_sampler()` from [`src/nanotron/data/samplers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/samplers.py) to ensure each rank processes a non-overlapping slice:

```python
from nanotron.data.samplers import get_sampler

dp_rank = parallel_context.dp_pg.rank()
dp_size = parallel_context.dp_pg.size()

sampler = get_sampler(
    train_dataset=my_dataset,
    dl_rank=dp_rank,
    dl_ranks_size=dp_size,
    seed=42,
    drop_last=True,
    shuffle=False,
    micro_batch_size=micro_batch,
)

```

### Assembling the DataLoader

Combine your custom collator with the distributed sampler using standard PyTorch DataLoader mechanics:

```python
import torch
from torch.utils.data import DataLoader

def make_my_dataloader(
    my_dataset,
    seq_len,
    micro_batch,
    parallel_context,
    input_pp_rank,
    output_pp_rank,
    seed=42,
    num_workers=4,
):
    sampler = get_sampler(
        train_dataset=my_dataset,
        dl_rank=parallel_context.dp_pg.rank(),
        dl_ranks_size=parallel_context.dp_pg.size(),
        seed=seed,
        drop_last=True,
        shuffle=False,
        micro_batch_size=micro_batch,
    )
    
    collator = MyCustomCollator(
        seq_len=seq_len,
        input_pp_rank=input_pp_rank,
        output_pp_rank=output_pp_rank,
        parallel_context=parallel_context,
    )
    
    return DataLoader(
        my_dataset,
        batch_size=micro_batch,
        sampler=sampler,
        collate_fn=collator,
        num_workers=num_workers,
        pin_memory=True,
    )

```

This pattern follows `get_train_dataloader()` in [`src/nanotron/data/dataloader.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/dataloader.py) while substituting your custom collator logic.

## Integrating with NanotronTrainer

The `NanotronTrainer` class in [`src/nanotron/trainer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/trainer.py) accepts a `train_dataloader` argument and stores it as an instance attribute. You can inject your custom DataLoader either during initialization or by direct assignment:

```python
from nanotron.trainer import NanotronTrainer
from nanotron.config import Config

cfg = Config.from_yaml("path/to/your/config.yaml")
trainer = NanotronTrainer(cfg)

custom_loader = make_my_dataloader(
    my_dataset,
    seq_len=cfg.model.sequence_length,
    micro_batch=cfg.training.micro_batch_size,
    parallel_context=trainer.parallel_context,
    input_pp_rank=0,
    output_pp_rank=trainer.parallel_context.pp_pg.size() - 1,
)

trainer.train_dataloader = custom_loader
trainer.fit()

```

The training loop iterates over `self.train_dataloader` and expects the batch dictionary format described above. As long as your custom DataLoader honors the contract, the pipeline parallelism, tensor parallelism, and context parallelism mechanisms function correctly.

## Key Source Files for Custom DataLoader Implementation

- **[`src/nanotron/data/dataloader.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/dataloader.py)**: Contains `get_train_dataloader()` which demonstrates the canonical DataLoader construction pattern with collator and sampler injection.
- **[`src/nanotron/data/clm_collator.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/clm_collator.py)**: Reference implementations of `DataCollatorForCLM` and `DataCollatorForCLMWithPositionIds` showing the exact batch dictionary structure required.
- **[`src/nanotron/data/samplers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/samplers.py)**: Implements `get_sampler()` for distributed dataset sharding across data-parallel ranks.
- **[`src/nanotron/data/dataloader_builder.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/dataloader_builder.py)**: Low-level helper for constructing DataLoaders with `Nanoset` datasets, useful for understanding worker initialization.
- **[`src/nanotron/trainer.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/trainer.py)**: The `NanotronTrainer` class that consumes the DataLoader; inspect the training step to understand batch consumption.

## Summary

- **Custom DataLoaders are fully supported** in Nanotron provided they return the specific batch dictionary schema with `input_ids`, `input_mask`, `label_ids`, and `label_mask`.
- **TensorPointers are mandatory** on pipeline-parallel ranks that do not own the corresponding tensors to maintain distributed consistency.
- **Use `get_sampler()`** from [`src/nanotron/data/samplers.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/samplers.py) to ensure proper dataset sharding across data-parallel processes.
- **Override `trainer.train_dataloader`** to inject your custom implementation without modifying the core training loop.
- **Maintain micro-batch size and sequence length alignment** between your collator output and the model's expected input dimensions.

## Frequently Asked Questions

### What happens if I don't use TensorPointers in my custom collator?

If you return actual tensors instead of TensorPointers on ranks that don't own the data, the pipeline-parallel communication hooks in Nanotron will attempt to access tensors that don't exist on those ranks, resulting in distributed runtime errors or deadlock during the forward pass.

### Can I use any PyTorch Dataset with a custom Nanotron DataLoader?

Yes, you can use any `torch.utils.data.Dataset` or `datasets.Dataset` implementation. However, your collator must transform the raw dataset samples into Nanotron's expected batch format. The dataset must also be compatible with the distributed sampler to prevent data duplication across data-parallel ranks.

### Do I need to handle sequence packing and position IDs manually?

Only if your training configuration sets `use_position_ids=True`. In that case, your collator must additionally provide `position_ids` as either a `torch.Tensor` or `TensorPointer`, following the same ownership rules as the other fields. Reference `DataCollatorForCLMWithPositionIds` in [`src/nanotron/data/clm_collator.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/data/clm_collator.py) for the implementation pattern.

### Is there a performance penalty for using a custom DataLoader?

No significant overhead exists provided your custom collator maintains the same tensor shapes and dtypes as the default implementation. The critical factor is ensuring your DataLoader uses `pin_memory=True` and appropriate `num_workers` to keep the GPU fed, matching the performance characteristics of `get_train_dataloader()` in the standard pipeline.