pytorch/torchtitan

torchtitan – PyTorch‑native platform for training generative AI models

What it is – torchtitan is an open‑source library built by the PyTorch team that lets you train large language models (LLMs) such as Llama 3.1 at scale using only native PyTorch primitives. It bundles a collection of distributed‑training techniques (FSDP2, tensor‑parallel, pipeline‑parallel, context‑parallel, etc.) together with utilities for checkpointing, quantization, and logging, all exposed through a clean, extensible Python API.

Why it matters – Training modern LLMs (8 B to 405 B parameters) normally requires a complex stack of custom CUDA kernels, third‑party libraries, and heavyweight orchestration tools. torchtitan aims to reduce that complexity by providing a minimal‑code‑base that works directly with the latest PyTorch nightly releases. Researchers can experiment with new parallelism strategies without rewriting model code, and production teams can move from prototype to multi‑node training with the same library.

Core capabilities (as listed in the README)

Category	Features
Parallelism	• Multi‑dimensional composable parallelisms: FSDP2 (per‑parameter sharding), Tensor Parallel (including async TP), Pipeline Parallel (zero‑bubble), Context Parallel for ultra‑long sequences.
Memory & Speed tricks	• Meta‑device model initialization, selective/full activation checkpointing, distributed async checkpointing, BF16 optimizer states, Float8 and MXFP8 quantization (Blackwell GPUs).
Training workflow	• `torch.compile` support, gradient accumulation via token‑count flag, flexible warm‑up‑stable‑decay LR scheduler, built‑in data loader with C4 pre‑configurable, custom‑dataset plug‑in.
Observability	• TensorBoard / Weights & Biases metrics (loss, memory, throughput, MFU), structured per‑rank logging, profiling tools (CPU/GPU, Flight Recorder).
Interoperability	• Distributed checkpoint format (DCP) that can be loaded directly into `torchtune` for fine‑tuning, TorchFT integration, helper scripts for tokenizer download and checkpoint conversion.
Extensibility	• Well‑documented extension points, an `experiments` folder for rapid prototyping (e.g., TitanRL RL stack), and clear guidelines for adding new models.
Supported hardware	• NVIDIA GPUs (tested up to H100, 512‑GPU runs), AMD ROCm fork, and nightly‑compatible builds for any CUDA/ROCm version.

Typical workflow

Install – pip install torchtitan (or use the nightly PyTorch + nightly torchtitan wheels). The package works out‑of‑the‑box on a single node or across Slurm/ParallelCluster clusters.
Prepare data – Use the built‑in C4 loader or point the config at a custom dataset.
Pick a model – Llama 3.1 (8 B, 70 B, 405 B) is shipped; adding a new model follows the short guide in torchtitan/models/README.md.
Configure – All knobs are exposed via a Python config file and CLI flags (--module, --config). You can enable/disable any parallelism, checkpointing, or quantization option.
Launch – Run torchrun (or the provided run_train.sh) on 1‑N nodes. torchtitan automatically composes the requested parallelisms and starts training.
Monitor – Metrics appear in TensorBoard/W&B; logs include per‑rank traces for debugging.
Checkpoint & fine‑tune – Saved DCP checkpoints can be loaded directly with torchtune for downstream SFT or RLHF.

Getting started (quick example)

# 1️⃣ Install the nightly PyTorch + torchtitan wheels
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130

# 2️⃣ Download the Llama‑3.1 tokenizer
python scripts/download_hf_assets.py \
    --repo_id meta-llama/Llama-3.1-8B \
    --assets tokenizer \
    --hf_token <your‑HF‑token>

# 3️⃣ Launch an 8‑GPU run (single node)
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

The CONFIG=llama3_8b file lives in torchtitan/config/ and already enables FSDP2 + Tensor Parallel + activation checkpointing, etc.

Who should use it?

Researchers exploring new scaling laws, parallelism combinations, or quantization techniques – the library’s clean‑room implementation makes it easy to read and modify the underlying code.
Engineers building production‑grade LLM pre‑training pipelines who prefer to stay within the PyTorch ecosystem and avoid external proprietary stacks.
Students learning distributed training concepts; the experiments folder contains runnable demos (e.g., TitanRL) that illustrate how the same model code can be reused for RL training and inference.

Community & support

Forum – Dedicated PyTorch forum category (pytorch‑forum → distributed → torchtitan).
Paper – ICLR 2025 paper “TorchTitan: One‑stop PyTorch native solution for production ready LLM pre‑training” (arXiv 2410.06511).
Contributing – Guidelines for core changes and for adding experimental ideas are in CONTRIBUTING.md and torchtitan/experiments/README.md.
Third‑party forks – An AMD‑optimized fork exists (AMD‑AGI/torchtitan‑amd).

License

torchtitan is released under the BSD‑3‑Clause license. The code itself is permissively licensed, but users must respect the licenses of any external data or model weights they download (e.g., Meta‑Llama tokenizer and checkpoints).

Bottom line – torchtitan provides a production‑ready, PyTorch‑only stack for training LLMs at scale, with a strong emphasis on readability, extensibility, and up‑to‑date distributed training research. If you need to train or experiment with large generative models without pulling in a heavyweight external ecosystem, this is the library to start with.