Nanotron - Entraînement LLM parallèle 3D minimaliste
Nanotron (par Hugging Face) est un cadre minimaliste pour le préentraînement de grands modèles de langage avec parallélisme 3D — données, tenseurs et pipeline parallèles — implémentés en Python délibérément lisible. Son objectif est d’être suffisamment petit et modifiable pour réellement comprendre, tout en distribuant l’entraînement sur de nombreux GPUs. C’est un excellent choix quand vous voulez apprendre ou personnaliser la mécanique de l’entraînement distribué sans la complexité des plus grands cadres.
L’entraînement est gourmand en calcul ; commencez petit. Validez une config sur un petit modèle et quelques GPUs avant de mettre à l’échelle.
Installation
| Méthode | Commande |
|---|
| Depuis la source | git clone https://github.com/huggingface/nanotron && cd nanotron && pip install -e . |
| Flash attention | pip install flash-attn (recommandé) |
| Exigences | PyTorch + GPUs CUDA |
| Vérifier | exécuter un exemple de config sur 1–2 GPUs |
Parallélisme 3D
| Dimension | Divise | Clé de config |
|---|
| Parallélisme des données (DP) | Le lot sur les répliques | dp |
| Parallélisme des tenseurs (TP) | Les couches sur les GPUs | tp |
| Parallélisme des pipelines (PP) | Le modèle en étapes | pp |
Le produit dp × tp × pp doit égaler votre nombre total de GPUs.
Configuration (YAML)
L’entraînement Nanotron est défini par un fichier de config YAML décrivant le modèle, les données, le parallélisme et l’optimiseur.
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
Lancement de l’entraînement
# 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
| Clé | Sens |
|---|
run_train.py | Point d’entrée d’entraînement principal |
--config-file | Le fichier de config YAML d’entraînement |
pp_engine | Calendrier du pipeline (par ex. 1F1B, AFAB) |
micro_batch_size | Lot micro par étape |
batch_accumulation_per_replica | Accumulation de gradient |
Données
| Option | Note |
|---|
| Datasets prétokenisés | Pointez la config vers les données tokenisées |
| Datasets HF | Flux/charges via Hugging Face |
| Nanoset | Format d’ensemble de données efficace de Nanotron |
| Tokenizer | Défini via la config tokenizer |
Points de contrôle et reprise
| Aspect | Comment |
|---|
| Fréquence d’enregistrement | checkpoints.checkpoint_interval |
| Chemin d’enregistrement | checkpoints.checkpoints_path |
| Reprendre | --config-file + pointez vers un checkpoint sauvegardé |
| Sharder | Les checkpoints connaissent le parallélisme |
Surveillance
| Métrique | Sens |
|---|
| loss / lm_loss | Perte d’entraînement |
| tokens/s | Débit |
| grad norm | Magnitude du gradient (stabilité) |
| Logging | Intégration W&B / TensorBoard |
Nanotron vs autres cadres d’entraînement
| Aspect | Nanotron | TorchTitan | Megatron-LM |
|---|
| Philosophie | Minimal, modifiable | Référence native PyTorch | Riche en fonctionnalités, à grande échelle max |
| Parallélisme | 3D (dp/tp/pp) | 4D (ajoute contexte) | 3D+ extensive |
| Lisibilité | Élevée (apprendre à partir de cela) | Élevée | Inférieure (grande base de code) |
| Meilleur pour | Comprendre/personnaliser 3D | Mise à l’échelle native PyTorch | Formation à grande échelle frontière |
Ressources
