TorchTitan - PyTorch-Native LLM Pretraining Cheatsheet
TorchTitan is an open-source, PyTorch-native platform for large-scale generative-model (LLM) pretraining. It is a clean-room reference implementation of PyTorch’s latest distributed-training techniques, composing 4D parallelism — data, tensor, pipeline, and context parallel — in a modular way, on top of FSDP2, torch.compile, and distributed checkpointing, with elastic scaling. It is the go-to starting point for teams that want to pretrain models using native PyTorch scaling rather than a third-party stack.
Distributed pretraining is resource-intensive. Validate configs on a small model/scale before committing a large multi-node run.
Installation
| Method | Command |
|---|
| From source | git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e . |
| Requirements | Recent PyTorch (nightly often recommended), CUDA GPUs |
| Datasets | Uses HF datasets / tokenizers |
| Verify | run a debug config on 1 GPU |
Parallelism Dimensions (4D)
| Dimension | Splits |
|---|
| Data Parallel (FSDP2) | The batch, sharding params/optimizer/grads |
| Tensor Parallel | Individual layers across GPUs |
| Pipeline Parallel | The model into stages across GPUs |
| Context Parallel | Long sequences across GPUs |
These compose: e.g. dp × tp × pp × cp maps a model onto a large GPU mesh.
Configuration
TorchTitan is driven by TOML config files plus command-line overrides.
# Launch training with a config (multi-GPU via torchrun)
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh
| Config section | Controls |
|---|
[model] | Architecture, size, flavor |
[training] | Batch size, steps, seq length, LR |
[parallelism] | data_parallel_shard_degree, tensor_parallel_degree, pipeline_parallel_degree, context_parallel_degree |
[optimizer] | Optimizer + hyperparameters |
[checkpoint] | Distributed checkpoint frequency/paths |
[activation_checkpoint] | Recomputation policy |
Key Features
| Feature | Benefit |
|---|
| FSDP2 | Modern fully-sharded data parallel |
torch.compile | Graph compilation for speed |
| Async distributed checkpointing | Save/resume huge runs efficiently |
| Float8 / low precision | Faster training on supported hardware |
| Elastic scaling | Adapt to changing GPU availability |
| Meta-device init | Build huge models without OOM at init |
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 | Example |
|---|
| 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 | Minimalistic 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
