Overview
TorchTitan is a PyTorch-native platform from the PyTorch team for pretraining and large-scale training of generative AI models such as Llama 3.1. It is built as a minimal, clean-room implementation of PyTorch's native scaling techniques, so the model code stays readable while multi-dimensional parallelism is applied around it.
It is aimed at researchers and engineers who want to train large models across many GPUs and experiment with new architectures or distributed-training techniques. Rather than hiding the internals, it keeps a small codebase with reusable, swappable components and documented extension points you can build on.
Within the efficient-training space, TorchTitan's focus is combining FSDP2, tensor parallel, pipeline parallel, and context parallel in a single composable stack, alongside features like torch.compile, Float8 training, distributed checkpointing, and activation checkpointing.
What it does
- Composable multi-dimensional parallelism: FSDP2 per-parameter sharding, Tensor Parallel (incl. async TP), Pipeline Parallel, and Context Parallel
- torch.compile support plus Float8 and MXFP8 (Blackwell) training for dense and MoE models
- Distributed checkpointing, including async checkpointing and checkpoints loadable in torchtune for fine-tuning
- Per-op selective and full activation checkpointing, plus meta-device initialization for large models
- Checkpointable data loading with the C4 dataset preconfigured (144M entries) and support for custom datasets
- Metrics for loss, GPU memory, throughput, TFLOPs, and MFU logged to TensorBoard or Weights & Biases
Getting started
Install TorchTitan from source (recommended to pair with a recent PyTorch nightly), download a tokenizer, then launch a training run.
Install from source
Clone the repo and install requirements. The README also recommends installing the nightly torchdata.
git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt
pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpuDownload a tokenizer
Fetch the tokenizer assets from Hugging Face (supply your own HF token).
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...Launch a training run
Run the provided launch script, selecting a model module and config.
MODULE=llama3 CONFIG=llama3_8b ./run_train.shCommands and code are distilled from the project's own documentation — always check the official repo for the latest.
When to use it
- Pretraining Llama 3.1 LLMs of various sizes across multiple GPUs or nodes
- Experimenting with combinations of FSDP, tensor, pipeline, and context parallelism on your own model
- Training long-context models using context parallel for very long sequence lengths
- Producing distributed checkpoints that can later be loaded into torchtune for fine-tuning
How TorchTitan compares
TorchTitan alongside other open-source efficient training tools AI/TLDR tracks, ranked by GitHub stars.
| Tool | Stars | What it does |
|---|---|---|
| DeepSpeed | ★ 42.8k | A deep learning optimization library whose ZeRO memory partitioning and offloading let you train very large models across many GPUs. |
| Megatron-LM | ★ 17.1k | NVIDIA's library for training large transformer models at scale using tensor, pipeline, and sequence parallelism. |
| Accelerate | ★ 9.8k | A library that runs the same PyTorch training code across CPUs, multiple GPUs, and TPUs while handling mixed precision, FSDP, and DeepSpeed. |
| Liger-Kernel | ★ 6.5k | A set of fused Triton kernels for common LLM layers that raises training throughput and lowers memory use as a drop-in replacement. |
| TorchTitan | ★ 5.5k | PyTorch-native platform for large-scale generative AI model training |
| Nanotron | ★ 2.8k | Hugging Face's minimal library for pre-training LLMs with 3D parallelism, designed to be readable and easy to modify. |