Nanotron - Minimalistic 3D-Parallel LLM Training Cheatsheet
Nanotron (by Hugging Face) is a minimalistic framework for pretraining large language models with 3D parallelism — data, tensor, and pipeline parallel — implemented in deliberately readable Python. Its goal is to be small and hackable enough to actually understand, while still distributing training across many GPUs. It is a strong choice when you want to learn or customize the mechanics of distributed pretraining without the complexity of the largest frameworks.
Pretraining is compute-heavy; start small. Validate a config on a tiny model and a couple of GPUs before scaling.
Installation
| Method | Command |
|---|
| From source | git clone https://github.com/huggingface/nanotron && cd nanotron && pip install -e . |
| Flash attention | pip install flash-attn (recommended) |
| Requirements | PyTorch + CUDA GPUs |
| Verify | run an example config on 1–2 GPUs |
3D Parallelism
| Dimension | Splits | Config key |
|---|
| Data Parallel (DP) | The batch across replicas | dp |
| Tensor Parallel (TP) | Layers across GPUs | tp |
| Pipeline Parallel (PP) | The model into stages | pp |
The product dp × tp × pp must equal your total GPU count.
Configuration (YAML)
Nanotron training is defined by a YAML config describing the model, data, parallelism, and optimizer.
parallelism:
dp: 2
tp: 2
pp: 2
pp_engine: 1f1b # pipeline schedule
model:
model_config:
hidden_size: 2048
num_hidden_layers: 24
num_attention_heads: 16
optimizer:
learning_rate_scheduler:
learning_rate: 3.0e-4
optimizer_factory:
adam_beta1: 0.9
adam_beta2: 0.95
tokens:
sequence_length: 2048
micro_batch_size: 4
batch_accumulation_per_replica: 8
Launching Training
# Single node, multiple GPUs
torchrun --nproc_per_node=8 run_train.py --config-file config.yaml
# Generate a starter config
python examples/config_tiny_llama.py
| Key | Meaning |
|---|
run_train.py | Main training entrypoint |
--config-file | The YAML training config |
pp_engine | Pipeline schedule (e.g. 1F1B, AFAB) |
micro_batch_size | Per-step micro batch |
batch_accumulation_per_replica | Gradient accumulation |
Data
| Option | Note |
|---|
| Pretokenized datasets | Point config at tokenized data |
| HF datasets | Stream/loads via Hugging Face |
| Nanoset | Nanotron’s efficient dataset format |
| Tokenizer | Set via the tokenizer config |
Checkpoints & Resume
| Aspect | How |
|---|
| Save frequency | checkpoints.checkpoint_interval |
| Save path | checkpoints.checkpoints_path |
| Resume | --config-file + point at a saved checkpoint |
| Sharded | Checkpoints are parallelism-aware |
Monitoring
| Metric | Meaning |
|---|
| loss / lm_loss | Training loss |
| tokens/s | Throughput |
| grad norm | Gradient magnitude (stability) |
| Logging | W&B / TensorBoard integration |
Nanotron vs Other Training Frameworks
| Aspect | Nanotron | TorchTitan | Megatron-LM |
|---|
| Philosophy | Minimal, hackable | PyTorch-native reference | Feature-rich, max-scale |
| Parallelism | 3D (dp/tp/pp) | 4D (adds context) | 3D+ extensive |
| Readability | High (learn from it) | High | Lower (large codebase) |
| Best for | Understanding/customizing 3D | Native PyTorch scaling | Frontier-scale training |
Resources
