Every serious lab eventually writes the same three files: a model, a sharding plan, and a training loop that quietly accretes hundreds of special cases for every parallelism combination. The PyTorch team looked at that pattern and decided to publish the reference version instead. torchtitan on GitHub is a clean-room, PyTorch-native platform for pretraining and post-training generative AI models, written from scratch so that FSDP2, tensor parallelism, multi-stage 1F1B schedules, context parallelism, and expert parallelism all stack into N-D mesh layouts without forcing you to rewrite the model. The project ships an ICLR 2025 paper, a v0.2.2 release, and verified recipes that have been measured on up to 512 GPUs.

What makes the codebase special is not a single clever trick but an architectural stance: configuration is a tree of dataclasses, models declare how they shard, and every run is a readable Python recipe you can copy and modify. In this post we walk the actual source - the engine that runs microbatches, the context that builds device meshes, the sharding placements, the quantized linear layers, and the 2026-era TitanRL stack that reuses the same model definitions inside a vLLM generator. As always in this series, what follows is an educational tour of published source code.

Overview of the torchtitan architecture: entry and config feed the Trainer, which drives the shared TrainingEngine over mesh-resolved parallelism machinery, common model definitions, and reusable components.

Reading the overview from left to right:

• torchtitan/training_engine.py is the heart at 872 lines: the TrainingEngine class starts at line 69, forward_backward at line 560, and optim_step at line 814.
• torchtitan/trainer.py is the application shell at 541 lines: Trainer at line 40 owns data policy, and train_step at line 332 assembles microbatch groups for one optimizer update.
• torchtitan/distributed/parallelism_context.py builds device meshes at 665 lines: ParallelismContext at line 100 and build_mesh at line 244 turn degree fields into DeviceMesh axes.
• torchtitan/distributed/fsdp.py applies FSDP2 at 499 lines: apply_fsdp_to_decoder at line 237 walks the model with per-parameter shard placements.
• torchtitan/models/common/decoder.py is the shared decoder at 401 lines: Decoder at line 68 with forward at line 228 and its own _apply_fsdp hook at line 132.
• torchtitan/models/common/moe/moe.py is the MoE toolbox at 1134 lines: TokenChoiceTopKRouter at line 171 and the MoE block at line 707.
• torchtitan/components/checkpointer/dcp.py persists state at 656 lines: CheckpointManager at line 89 with async dcp_save at line 252.
• torchtitan/config/transform/quantization.py rewrites linears at 392 lines: MXFP8LinearConverter at line 91 and NVFP4LinearConverter at line 277.

Why You Need This

If you have ever tried to bolt tensor parallelism onto an FSDP training script, you know the failure mode: two sharding systems fight over the same parameters, and the fix is usually a fork of the whole trainer. torchtitan attacks that at the root by routing everything through torch.distributed DeviceMesh and DTensor. The ParallelismContext turns plain integer degrees into a validated multi-axis mesh; the FSDP2 application, tensor-parallel placements, and stage splitting all read from that one mesh. Your model code keeps calling plain torch ops and stays byte-for-byte identical whether it runs on 1 GPU or 512.

The second reason is the recipe model of configuration. Instead of a YAML file with 90 keys whose semantics nobody remembers, a recipe is a Python function that returns a complete Trainer.Config - see how llama3_8b in the verified recipes folder sets loss, optimizer, dataloader, and remat policy in 35 readable lines. Changing a run means writing a function, diffing it in review, and importing it in CI. The repo even validates loss convergence across parallelism combinations, so the defaults you inherit are measured, not aspirational.

The third reason is that this is where PyTorch itself showcases frontier techniques first: MXFP8 and NVFP4 quantized training for Blackwell-scale hardware, distributed MoE with fused dispatch and combine, regional torch.compile and CUDA graphs, and the TitanRL post-training stack that shares model definitions between the trainer and a vLLM generator with a bitwise-reproducible batch-invariant mode. Reading this repo is effectively reading the PyTorch team’s opinionated answer to “how should large-scale training look in 2026”. Where the torchao quantization library we covered earlier supplies the kernels, torchtitan supplies the wiring that puts those kernels into a real training loop.

Detailed view: config loaders and recipes feed the engine core; the parallelism group shows mesh building, FSDP placement, remat and CUDA graph wrappers; the model zoo machinery and quantization/RL groups sit alongside.

Start where execution starts. torchtitan/train.py is only 72 lines: main at line 22 loads a ConfigLoader, builds the Trainer, and hands off to trainer.train(), which loads the latest checkpoint, loops should_continue_training, and saves per step. The shell script run_train.sh wraps torchrun with NGPU, MODULE, and CONFIG environment variables at lines 28 through 32, and even supports a fake backend mode at line 42 that dry-runs a 32-GPU layout on one device without real communication - a blessing for testing topology changes. The Trainer itself, at line 40 of the 541-line trainer.py, validates in post_init that you have not combined incompatible options (one guard at lines 86 through 96 rejects validation together with CUDA graphs on multi-stage schedules), then builds the dataloader, tokenizer, and engine. The interesting orchestration lives in train_step at line 332: it fills microbatch_groups for gradient_accumulation_steps, folds each microbatch’s loss and routing token counts into local tensors, and after the engine call reduces them across the data-parallel mesh with dist_sum so the logged loss is globally normalized rather than rank-local.

The engine is where the parallelism contract pays off. training_engine.py documents its own call flow in a docstring at lines 77 through 107, and forward_backward at line 560 implements exactly that: it zeroes gradients, converts global token counts to device tensors, optionally defers microbatch preprocessing until after CUDA graph capture via preprocess_microbatch_groups_lazily, and can route the whole call through an SDC replayer that detects silent data corruption by replaying the forward/backward. The _forward_backward_body at line 681 iterates accumulation groups, and each group goes to either _pp_forward_backward_microbatch_group at line 778, which packs argument, keyword, and target microbatches for multi-stage schedules, or _non_pp_forward_backward_microbatch at line 756 for the simple case. optim_step at line 814 then advances the optimizer, LR schedulers, and EMA inside the same optional CUDA graph capture. Nothing in this file mentions a specific model: it speaks only in model_parts and microbatches, which is exactly why one engine serves Llama, Qwen, DeepSeek, and FLUX alike.

The parallelism machinery is the most transferable code in the repo. parallelism_context.py defines ParallelismContext at line 100 with build_mesh at line 244, which unfolds the dp axis into replicate and shard sub-axes and validates the mesh constraint documented for expert parallelism in config/parallelism.py at line 174: the dense region (dp_shard times cp times tp) must cover the same ranks as the sparse region (edp_shard times ep). The ParallelismConfig dataclass at line 22 reads like a distributed-systems syllabus - fsdp_defer_gradient_reduction at line 61, fsdp_symm_mem_scope at line 64 for symmetric-memory all-gathers, context_parallel_degree at line 164 - and post_init at line 193 cross-checks schedule names against torch.distributed before you ever launch. Downstream, fsdp.py provides apply_fsdp_to_decoder at line 237 driven by linear_param_shard_placements at line 50, while context_parallel.py contributes load balancers like PTRRFlexAttentionCPLoadBalancer at line 97 that shard attention metadata across context ranks. Memory features complete the picture: activation_remat.py defines FullActivationRematConfig at line 167, and cuda_graph.py wraps forward, backward, and even the optimizer with CUDAGraphWrapper at line 296 after NUM_CUDA_GRAPH_WARMUP_STEPS at line 25 warms the capture.

The model zoo shows how little code a model needs once the base classes carry the distributed weight. decoder.py implements Decoder at line 68 with init_states, forward at line 228, preprocess_inputs at line 272, and _get_attention_metadata at line 378; models only supply block structure and sharding hints. The attention catalog in attention/attention.py spans 953 lines: GQAttention at line 185 for grouped-query heads, VarlenAttentionMetadata at line 288 with VarlenInnerAttention at line 314 for packed variable-length sequences, FlexInnerAttention at line 439 wrapping FlexAttention (the kernels connect back to the FlashAttention lineage we toured earlier), and sliding-window plus document mask-mod builders near lines 745 through 833. The MoE stack in moe/moe.py is 1134 lines of router engineering - TokenChoiceTopKRouter at line 171, QuantileBalancedTopKRouter at line 356, and two auxiliary losses at lines 511 and 609 - while dist_moe/runtime.py adds DistMoeRuntime at line 146 that plans fused dispatch, expert compute, and combine activation slots across stage boundaries on Blackwell. A concrete model is then tiny: llama3/model.py is 112 lines where Llama3TransformerBlock at line 28 annotates residual adds with torch_remat regions at lines 63 through 70 so backward replay can skip what needs no persisting, and state_dict_adapter.py handles HF checkpoint conversion with a qkv reverse-permutation at line 66 before from_hf at line 128.

The components layer and the RL stack round out the tour. components/loss.py is 848 lines with real autograd Functions - _LossParallelCrossEntropy at line 67 and _VocabParallelEntropy at line 222 compute vocab-parallel losses without gathering logits, and ChunkedLossWrapper at line 504 tiles the cross-entropy over sequence chunks to spare memory. components/optim/optimizer.py contributes AdamW at line 148, the Muon-style DistMuon at line 186, and OptimizersContainer at line 225 that installs BF16 moment hooks registered at line 69. Checkpointing is split between checkpointer/base.py (BaseCheckpointManager at line 198 with a purge_thread at line 44 enforcing retention policy) and checkpointer/dcp.py, whose dcp_save at line 252 writes torch.distributed.checkpoint state with an async staging path guarded by _maybe_wait_for_staging at line 523 - saves overlap with training, which matters at terabyte scale. Metrics live in observability/metrics.py with MetricsProcessor at line 305 reporting TFLOPs and MFU per step. Quantization converts the model tree before training starts: the converters in quantization.py swap Linear modules for quantization/mxfp8/linear.py MXFP8Linear at line 372 (autograd Function at line 74) or quantization/nvfp4/linear.py NVFP4Linear at line 330 with its random Hadamard sign vector at line 460. Finally TitanRL reuses everything: rl/train.py provisions trainer and generator hosts with PerHostProvisioner at line 99 and spawn_proc_mesh at line 230, while rollouts come from rl/generate.py - same model definitions, bitwise-identical numerics when batch-invariant mode is on.

From Install to Training

The quickest path is the README’s debug model, which needs no download because it uses a tokenizer bundled with the tests:

git clone https://github.com/pytorch/torchtitan && cd torchtitan
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu132
pip install -r requirements.txt
NGPU=1 ./run_train.sh        # Llama 3 debug model, 10 steps
MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8b ./run_train.sh

To make a run yours, write a recipe function instead of editing flags - the config stays a reviewable Python object:

# my_recipes.py
from torchtitan_recipes.models.llama3 import llama3_8b

def llama3_8b_tp2():
    config = llama3_8b()
    config.parallelism.tensor_parallel_degree = 2
    return config

Then launch with MODULE=my_recipes CONFIG=llama3_8b_tp2 ./run_train.sh and watch the mesh come up.

Honest limits: torchtitan targets recent PyTorch nightlies and modern accelerators, so it is not the tool for a laptop CPU or an old CUDA card; the model catalog, while growing fast (Llama 3, Qwen3 family, DeepSeek V3/V4, gpt-oss, Kimi, FLUX), is deliberately narrower than the whole Hugging Face zoo, and adding a new architecture means following the documented model guide rather than loading arbitrary checkpoints. Several advanced combinations carry real restrictions - validation is rejected together with CUDA graphs on multi-stage schedules, stage-parallel runs require schedule-aware loss visibility handling, and the quantized paths assume torchao nightlies matching your torch build. It is a training platform, not an inference server, and it assumes you are comfortable with torchrun, NCCL, and reading logs. None of that dims the core achievement: a readable, clean-room reference for how N-D parallel training should be assembled.

Enjoyed this post? Never miss out on future posts by following us