Nanotron is a library for pretraining transformer models, offering a simple, performant, and scalable API for custom datasets. It targets researchers and engineers building large language models, enabling efficient training through advanced parallelism techniques.

How It Works

Nanotron implements 3D parallelism (Data Parallelism, Tensor Parallelism, Pipeline Parallelism) to distribute model training across multiple GPUs and nodes. It supports expert parallelism for Mixture-of-Experts (MoE) models and includes optimized scheduling (AFAB, 1F1B) for pipeline parallelism. The library provides explicit APIs for TP and PP, facilitating debugging and customization, along with ZeRO-1 optimizer and FP32 gradient accumulation for memory efficiency.

Quick Start & Requirements

• Installation:

  uv venv nanotron --python 3.11 && source nanotron/bin/activate
  uv pip install torch --index-url https://download.pytorch.org/whl/cu124
  uv pip install -e .
  uv pip install datasets transformers datatrove[io] numba wandb ninja triton "flash-attn>=2.5.0" --no-build-isolation
  huggingface-cli login
  wandb login
  git-lfs --version
  

• Prerequisites: Python 3.11, PyTorch with CUDA 12.4, Git LFS.
• Resources: Requires multiple GPUs (e.g., 8 x H100s for the tiny Llama example).
• Docs: Ultrascale Playbook, Your First Training

Highlighted Details

• Supports 3D parallelism (DP+TP+PP) and expert parallelism for MoEs.
• Includes AFAB and 1F1B schedules for pipeline parallelism.
• Features ZeRO-1 optimizer, FP32 gradient accumulation, and parameter tying/sharding.
• Offers spectral µTransfer parametrization for scaling neural networks.

Maintenance & Community

• Actively developed by Hugging Face.
• Examples cover custom dataloaders, Mamba, MoE, and µTransfer.
• Roadmap includes FP8 training, ZeRO-3, torch.compile, and ring attention.

Licensing & Compatibility

• Licensed under Apache 2.0.
• Permissive license suitable for commercial use and integration with closed-source projects.

Limitations & Caveats

• FP8 training and ZeRO-3 (FSDP) are on the roadmap, not yet implemented.
• torch.compile support is also planned for future releases.