torchtitan – PyTorch‑native platform for large‑scale generative‑AI training

What it is

torchtitan is an open‑source library from the PyTorch team that lets you train, fine‑tune, and run reinforcement‑learning (RL) on very large language‑model (LLM) and multimodal models entirely in PyTorch. The codebase is deliberately small and readable, and it expresses the whole training pipeline as a Python “recipe” – a function that returns a configuration object. By swapping the recipe you can change model size, data, parallelism strategy, or even switch from pre‑training to RL without rewriting the model.

Core ideas & why they matter

Idea	What it does	Why it’s useful
Composable parallelism	Models declare how they should be sharded; the library stacks Fully‑Sharded Data Parallel (FSDP2), Tensor Parallel (TP), Pipeline Parallel (PP), Context Parallel (CP) and Expert Parallel (EP) via a DeviceMesh and DTensor.	Lets researchers experiment with any combination of parallelism on the same code, scaling from a single GPU to hundreds without custom kernels.
One‑stack from pre‑training to RL	The same model definition and kernels are used for classic language‑model pre‑training, supervised fine‑tuning (SFT), and the TitanRL reinforcement‑learning loop that also powers vLLM generation.	Guarantees bit‑wise reproducibility between trainer and inference, simplifying debugging of RL‑based alignment or policy‑gradient work.
Frontier model support	Ready‑made recipes for Llama 3, Qwen 3, DeepSeek V3, FLUX, Kimi, Muse Glimmer, etc., including low‑precision formats (MXFP8, NVFP4) and MoE (Dist‑MoE).	Researchers can start experiments on state‑of‑the‑art models without re‑implementing architecture or quantisation tricks.
Python‑first configuration	A recipe is just a Python function that returns a config object; you edit the function instead of passing a long list of CLI flags.	Makes experiments reproducible and easy to version‑control; you can share a recipe file with a colleague and they get the exact same run.
Extensive observability	Built‑in metrics (loss, TFLOPs, MFU), TensorBoard/W&B logging, profiling tools, deterministic mode, loss‑comparison scripts, and SDC replay.	Helps spot performance bottlenecks or numerical bugs early, which is critical at the scale of hundreds of GPUs.

What you can do with it today

• Train a new LLM – pick a recipe like llama3_8b and launch ./run_train.sh on 8 GPUs; the library handles sharding, checkpointing, and data loading automatically.
• Fine‑tune with LoRA or SFT – the same recipe system lets you add LoRA adapters or chat‑dataset fine‑tuning steps.
• Run RL‑based alignment – use the TitanRL stack to run PPO‑style or newer GRPO/DAPO losses, re‑using the exact model weights that were trained.
• Experiment with low‑precision training – enable MXFP8 or NVFP4 via a config flag; the library ships the necessary torchao kernels.
• Scale to hundreds of GPUs – documented benchmarks show up to 512‑GPU runs on H100s, with automatic distributed checkpointing and async saves.

Who should use it

• Research labs that need a clean, hackable codebase for exploring new parallelism or RL ideas.
• Engineers building production‑grade LLM training pipelines who want to stay within the PyTorch ecosystem (no separate C++/CUDA frameworks).
• Students learning large‑scale training – the recipes are readable examples of how to combine FSDP, TP, and MoE.

Getting started (quick summary)

1. Install – clone the repo, install a PyTorch nightly and the torchtitan package (or use the pre‑built pip/conda wheels).
2. Run a debug job – NGPU=1 ./run_train.sh trains a tiny Llama‑3 debug model for 10 steps.
3. Train a real model – download the tokenizer from Hugging Face, set MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8b, and launch on 8 GPUs.
4. Customize – write your own recipe function that tweaks parallelism degrees, data mix, or adds LoRA.

Documentation & community

• Docs are split into focused sections (config, models, distributed, data, checkpointing, RL, debugging, benchmarks). Each folder contains a README.md with usage examples.
• CI badges show active testing on CPU, NVIDIA (including H100/B200), AMD ROCm, and RL workloads.
• Contributing – guidelines encourage new experiments in the experiments/ folder and core improvements via pull requests.
• Paper & citation – an ICLR‑2025 paper details the design choices; cite it if you publish work built on torchtitan.

---

All information above is taken directly from the repository’s README; no external assumptions have been added.