TorchTitan - Préentraînement LLM natif PyTorch
TorchTitan est une plateforme open-source native PyTorch pour le préentraînement de modèles génératifs (LLM) à grande échelle. C’est une implémentation de référence épurée des dernières techniques d’entraînement distribué de PyTorch, composant le parallélisme 4D — parallélisme des données, tenseurs, pipeline et contexte — de manière modulaire, sur la base de FSDP2, torch.compile et checkpointing distribué, avec mise à l’échelle élastique. C’est le point de départ idéal pour les équipes qui veulent préentraîner des modèles en utilisant la mise à l’échelle native PyTorch plutôt qu’une pile tierce.
L’entraînement distribué est intensif en ressources. Validez les configs sur un petit modèle/échelle avant de vous engager dans une grande exécution multi-nœud.
Installation
| Méthode | Commande |
|---|
| Depuis la source | git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e . |
| Exigences | PyTorch récent (nightly souvent recommandé), GPUs CUDA |
| Datasets | Utilise les datasets HF / tokenizers |
| Vérifier | exécuter une config de débogage sur 1 GPU |
Dimensions du parallélisme (4D)
| Dimension | Divise |
|---|
| Parallélisme des données (FSDP2) | Le lot, sharding des params/optimiseur/grads |
| Parallélisme des tenseurs | Les couches individuelles sur les GPUs |
| Parallélisme des pipelines | Le modèle en étapes sur les GPUs |
| Parallélisme du contexte | Les longues séquences sur les GPUs |
Celles-ci se composent : par ex. dp × tp × pp × cp mappe un modèle sur une grande grille GPU.
Configuration
TorchTitan est piloté par des fichiers de config TOML plus des remplacements en ligne de commande.
# Launch training with a config (multi-GPU via torchrun)
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh
| Section config | Contrôle |
|---|
[model] | Architecture, taille, saveur |
[training] | Taille du lot, étapes, longueur seq, LR |
[parallelism] | data_parallel_shard_degree, tensor_parallel_degree, pipeline_parallel_degree, context_parallel_degree |
[optimizer] | Optimiseur + hyperparamètres |
[checkpoint] | Fréquence/chemins des points de contrôle distribués |
[activation_checkpoint] | Politique de recomputation |
Caractéristiques clés
| Caractéristique | Avantage |
|---|
| FSDP2 | Parallélisme de données entièrement sharder moderne |
torch.compile | Compilation de graphe pour la vitesse |
| Checkpointing distribué asynchrone | Enregistrer/reprendre efficacement les grandes exécutions |
| Float8 / basse précision | Formation plus rapide sur le matériel pris en charge |
| Mise à l’échelle élastique | S’adapter à la disponibilité des GPU changeante |
| Initialisation du dispositif Meta | Construire d’énormes modèles sans OOM lors de l’init |
Lancement multi-nœud
# torchrun across nodes (conceptual)
torchrun --nnodes=4 --nproc_per_node=8 \
--rdzv_backend=c10d --rdzv_endpoint=$MASTER:29500 \
-m torchtitan.train --job.config_file llama3_70b.toml
| Remplacer | Exemple |
|---|
| Degrés parallèles | --parallelism.tensor_parallel_degree 8 |
| Étapes | --training.steps 10000 |
| Longueur seq | --training.seq_len 8192 |
| Compiler | --training.compile |
Surveillance et points de contrôle
| Aspect | Comment |
|---|
| Logging | Intégration TensorBoard / W&B |
| Métriques | Loss, tokens/sec, MFU (utilisation du modèle FLOPs) |
| Points de contrôle | Checkpoint distribué (DCP) save/load |
| Reprendre | Pointez le config vers une étape de checkpoint |
TorchTitan vs autres cadres d’entraînement
| Aspect | TorchTitan | Megatron-LM | Nanotron |
|---|
| Origine | Équipe PyTorch | NVIDIA | Hugging Face |
| Style | Natif PyTorch, composable | Riche en fonctionnalités, mature | Minimaliste 3D |
| Parallélisme | 4D (dp/tp/pp/cp) | 3D+ (extensive) | 3D |
| Meilleur pour | Mise à l’échelle native PyTorch de référence | Modèles MoE/dense à grande échelle | Formation 3D petite et modifiable |
Ressources
