v0.4.0

TorchTitan 0.4.0

TorchTitan 0.4.0 aligns with PyTorch 2.15, expands model coverage, improves large-scale MoE training, and unifies core training and TitanRL on a shared engine.

These highlights cover changes since 0.3.0. For broader updates and benchmarks, see TorchTitan Update: New SOTA Models, Performance at Scale, TitanRL.

New model implementations

• Muse Spark Open: Adds the 816B-A42B model, including text and vision/audio configurations. Multimodal CUDA graphs remain unsupported. (#5420)
• Kimi K3: Hybrid KDA/MLA attention, LatentMoE, and vision support, with tensor/sequence and expert parallelism plus a dedicated pipeline implementation for attention residuals. (#4025, #4312)
• DeepSeek V4: Initial Flash/Pro support with sparse attention, indexer loss, optional multi-token prediction, and debug recipes. Full-scale checkpoint and convergence validation remain outstanding. (#3634, #4835)

Training performance and memory

• Dist-MoE integration: Fuses expert dispatch, computation, and combine. DeepSeek V3 BF16/MXFP8 recipes for NVIDIA SM100+ GPUs include fused MLA and SwiGLU. (#4541, #4542)
• Low-precision training: MXFP8 gains FSDP-managed prepared weights and fused weight-gradient accumulation across microbatches. Experimental NVFP4 support extends to grouped MoE linear layers. Both paths require supported Blackwell hardware and their respective optional dependencies. (#4203, #4911, #4821)
• Local compilation: model.local_compile_regions selects functions such as SwiGLU, RoPE, and gated RMSNorm for independent compilation and kernel fusion with torch.compile. (#4894, #4931, #4985, #5026)
• Faster router and LM-head layers: HiMidLoLinear uses BF16 matrix multiplications with FP32 outputs. Its backward pass combines two or three BF16 parts of each FP32 output gradient, improving accuracy over a single BF16 cast. (#4923)
• CUDA graphs and pipeline memory: Adds pipeline forward/backward capture, shared training/validation graph memory, and per-microbatch-group replay, plus controls for FSDP parameter residency and unshard lookahead. (#4493, #5241, #4592)
• Context parallelism: Adds Ulysses support for standard and variable-length attention, improves MLA communication, and adds native context-parallel paths for Qwen3.5 Gated DeltaNet and Kimi K3 KDA. (#4450, #4760, #5184)
• Activation rematerialization: torch_remat controls activation saving and recomputation by region. DeepEP dispatch preserves token order during recomputation, preventing incorrect MoE gradients. (#4497, #5115)
• FlexShard / DistMuon: Extends optimizer-compute layouts with variable-sized BlockShard partitions and BF16 redistribution of optimizer inputs and update directions. (#4840, #5279)

TitanRL

• TitanRL moves out of the experimental folder: Now in torchtitan.rl (formerly torchtitan.experiments.rl), it shares the core TrainingEngine and supports pipeline and context parallel training. (#4773, #3980, #5424)
• Extensible execution: ActorSystem separates service placement from RL logic, while GenerationEngine separates inference execution from request dispatch. Monarch and vLLM remain the provided defaults. (#5391, #5317)
• Multi-turn and agentic workloads: Optional session KV retention reuses cached prefixes across turns. Adds Verifiers and Terminal-Bench examples, plus a Muse Glimmer recipe for Search-R1. (#5220, #4356, #4942, #4145)

Data, checkpointing, and debugging

• Grain data pipeline: Composable dataset sources, mixing, packing, prefetching, and checkpointable iterator state replace the previous loader. Multi-turn text SFT adds renderer-based processing, tool metadata, and assistant-token loss masks. (#4088, #2769)
• Batch dimension folding: Flattens language-model batches from [B, L, D] to [T, D], enabling token-based batching with less SFT padding. (#4121)
• Checkpointing: Adds native Hugging Face safetensors loading via torch_checkpointing and parallel HF reads via DCP. Resume fixes initialize EMA when its state is absent and preserve configured BF16 Adam/AdamW moments. (#4277, #5087, #5141, #4615)
• Deterministic replay: An optional SDC check replays forward/backward execution, including gradient accumulation and pipeline microbatches, to detect mismatches. Deterministic execution is required. (#4333, #4334)
• GraphTrainer improvements: Adds full Inductor compilation, spmd_types tracing, and improved in-place weight-gradient accumulation to the experimental graph trainer. (#4222, #4325, #5208)

Compatibility and upgrade notes

• Recipe configuration: Use fully qualified Python recipe modules with --module and --config; move general CLI overrides into your recipe function. Rebuild model configurations when changing sequence or expert parallelism. (configuration guide)
• Removed features: Whole-TransformerBlock compilation in the standard trainer is replaced by model.local_compile_regions. Legacy Float8 training, MinimalAsyncEP, and the full_dtensor backend have been removed. (#4895, #4961, #4627, #4217)
• MXFP8 dependencies: Use a TorchAO nightly build for MXFP8 training in TorchTitan 0.4.0. TorchAO 0.18 does not include the required kernels.

Thank you to everyone who contributed code, testing, bug reports, and feedback!

Full Changelog: v0.3.0...v0.4.0