TorchTitan - PyTorch-Native LLM Pretraining Cheatsheet
TorchTitan ist eine Open-Source-PyTorch-native Plattform für großflächiges generatives Modell (LLM) Pretraining. Es ist eine saubere Referenzimplementierung der neuesten verteilten Trainingstechniken von PyTorch, komponiert 4D-Parallelität — Daten, Tensor, Pipeline und Kontext Parallel — auf modulare Weise auf Basis von FSDP2, torch.compile und verteiltem Checkpointing mit elastischer Skalierung. Es ist der Ausgangspunkt für Teams, die Modelle mit nativer PyTorch-Skalierung trainieren möchten, anstatt einen Drittanbieter-Stack zu verwenden.
Verteiltes Pretraining ist ressourcenintensiv. Validiere Konfigurationen auf einem kleinen Modell/Skala, bevor du einen großen Multi-Node-Lauf durchführst.
Installation
| Method | Command |
|---|
| From source | git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e . |
| Requirements | Aktuelle PyTorch (Nightly oft empfohlen), CUDA GPUs |
| Datasets | Uses HF datasets / tokenizers |
| Verify | run a debug config on 1 GPU |
Parallelism Dimensions (4D)
| Dimension | Splits |
|---|
| Data Parallel (FSDP2) | Den Batch über Replicas sharden, Parameter/Optimizer/Grads |
| Tensor Parallel | Einzelne Schichten über GPUs |
| Pipeline Parallel | Das Modell in Stufen über GPUs |
| Context Parallel | Lange Sequenzen über GPUs |
Diese komponieren: z.B. dp × tp × pp × cp mappt ein Modell auf ein großes GPU-Mesh.
Configuration
TorchTitan wird von TOML-Konfigurationsdateien plus Kommandozeilen-Überschreibungen angetrieben.
# Launch training with a config (multi-GPU via torchrun)
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh
| Config section | Steuert |
|---|
[model] | Architektur, Größe, Flavor |
[training] | Batch Size, Schritte, Seq-Länge, LR |
[parallelism] | data_parallel_shard_degree, tensor_parallel_degree, pipeline_parallel_degree, context_parallel_degree |
[optimizer] | Optimizer + Hyperparameter |
[checkpoint] | Verteilte Checkpoint-Häufigkeit/Pfade |
[activation_checkpoint] | Recomputation Policy |
Key Features
| Feature | Vorteil |
|---|
| FSDP2 | Modernes Fully-Sharded Data Parallel |
torch.compile | Graph-Kompilierung für Geschwindigkeit |
| Async distributed checkpointing | Speichere/fortsetze riesige Läufe effizient |
| Float8 / low precision | Schnelleres Training auf unterstützter Hardware |
| Elastic scaling | Passe dich wechselnder GPU-Verfügbarkeit an |
| Meta-device init | Baue riesige Modelle ohne OOM beim Initialisieren |
Launching Multi-Node
# 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
| Override | Beispiel |
|---|
| Parallel degrees | --parallelism.tensor_parallel_degree 8 |
| Steps | --training.steps 10000 |
| Seq length | --training.seq_len 8192 |
| Compile | --training.compile |
Monitoring & Checkpoints
| Aspect | How |
|---|
| Logging | TensorBoard / W&B Integration |
| Metrics | Loss, tokens/sec, MFU (Model FLOPs Utilization) |
| Checkpoints | Distributed Checkpoint (DCP) Save/Load |
| Resume | Point the config at a checkpoint step |
TorchTitan vs Other Training Frameworks
| Aspect | TorchTitan | Megatron-LM | Nanotron |
|---|
| Origin | PyTorch Team | NVIDIA | Hugging Face |
| Style | PyTorch-Native, Composable | Feature-Rich, Mature | Minimalistisch 3D |
| Parallelism | 4D (dp/tp/pp/cp) | 3D+ (Extensive) | 3D |
| Best for | Native PyTorch Reference Scaling | Max-Scale MoE/Dense | Small, Hackable 3D Training |
Resources
