torchtitan

torchtitan is a PyTorch-native, open-source platform from the PyTorch project for large-scale training of generative AI models. It takes a model from pretraining through supervised fine-tuning to reinforcement learning in one readable codebase. The latest release listed is v0.3.0, and the README notes the project is under active development and recommends a recent PyTorch nightly for the latest features.

What It Is

torchtitan is a clean-room implementation of large-scale training that composes FSDP2, tensor, pipeline, context, and expert parallelism without rewriting the model. Models declare how they shard, and the parallelisms stack into N-D layouts through DeviceMesh and DTensor. Each run is defined by a Python recipe, a function that returns a complete config, so changing a run means writing or copying a function rather than passing flags.

Models and Training Techniques

The README lists support for Llama 3, Qwen3 and later Qwen versions, DeepSeek V3 and V4, gpt-oss, Kimi models, Muse Glimmer, and FLUX. Documented techniques include activation checkpointing, regional torch.compile and CUDA graphs, MXFP8 and NVFP4 low-precision training, FlashAttention, FlexAttention, BF16 optimizer states, and FSDP CPU offload. Training covers pretraining, SFT on chat datasets, and LoRA, with a Grain data pipeline for packing, mixing, and checkpointable state.

TitanRL for Post-Training

TitanRL reuses the same model definitions and kernels in the trainer and in vLLM generation, with a batch-invariant mode for bitwise trainer/generator parity. It runs async rollouts orchestrated with Monarch, syncs weights over TorchStore, and supports GRPO and DAPO losses. Examples include DAPO Math, Search-R1, and Verifiers.

Checkpointing and Debugging

Distributed checkpointing supports async saves and direct loading and saving of Hugging Face safetensors, with conversion scripts between HF and DCP formats. Observability covers loss, memory, throughput, TFLOPs, and MFU in TensorBoard or W&B, plus profiling, memory snapshots, and Flight Recorder. Fake-backend dry runs of large layouts on a single GPU, deterministic mode, and loss comparison across commits are also documented.

Setup Path

Install from source with a PyTorch nightly on Python 3.11+, or use nightly builds, pip, or conda. A single-GPU debug Llama 3 run needs no tokenizer download, and a multi-node Slurm script is provided.