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

What's changed

Compared with v0.2.2, TorchTitan 0.3.0 adds typed Python configuration, expands model and declarative parallelism support, and enhances MoE execution, graph-native training, training efficiency, SFT, and RL.

Configuration architecture and extensibility

Recipes now support Python inheritance and typed CLI overrides. An opt-in override registry also lets external packages replace selected model or training components without modifying TorchTitan core.

Key PRs:

• Python configuration refactor (#2386)
• External override registry (#3396)

Models and architecture

This release adds Qwen3.5 with variable-length attention, Kimi K2.7 with DistributedMuon, Muse Glimmer 30B, DeepSeek V3 MTP and fused MLA, and Flux compilation with MXFP8.

Key PRs:

• Qwen3.5 (#3371)
• Qwen3.5 variable-length attention (#3801)
• Kimi K2.7 (#3532)
• Kimi K2.7 DistributedMuon overlap (#4060)
• Muse Glimmer 30B (#4106)
• Muse Glimmer FA3 and vLLM support (#4126)
• DeepSeek V3 MTP (#3392)
• DeepSeek V3 fused MLA (#4134)
• Flux compile and MXFP8 (#2579)

Declarative SPMD

spmd_types is now the default SPMD backend, providing declarative state and activation sharding across Llama 3, Qwen3, DeepSeek V3, GPT-OSS, Flux, Qwen3.5, Kimi K2.7, Muse Glimmer, and the Transformers modeling backend.

Some GraphTrainer configurations continue to use the legacy partial_dtensor backend, where model-parallel axes use DTensor while FSDP is managed separately.

Read the SPMD types integration devlog for more details.

Key PRs:

• State sharding and redistribution (#3587)
• SPMD infrastructure (#3641)
• Backend configuration (#3466)
• Llama 3 (#3763)
• Qwen3 (#3655)
• DeepSeek V3 (#3673)
• GPT-OSS (#3690)
• Flux (#3823)
• Qwen3.5 (#3895)
• Kimi K2.7 (#4079)
• Muse Glimmer (#4161)
• Default backend (#4085)
• Transformers modeling backend (#4235)

MoE execution and expert parallelism

MoE is reorganized around a unified token-dispatcher interface that separates routing, token movement, and expert computation. Standard, MinimalAsyncEP, DeepEP, and HybridEP backends support communication overlap, CUDA graphs, and composition with TP, CP, PP, activation checkpointing, and compilation.

Key PRs:

• Token dispatcher (#2842)
• Dispatcher decomposition (#3389)
• Unified dispatcher API (#3970)
• SPMD sparse-mesh transitions (#3654)
• MinimalAsyncEP (#3561)
• MinimalAsyncEP composition (#4089)
• HybridEP compilation (#3360)
• DeepEP v2 (#3808)

GraphTrainer and GraphPP

GraphTrainer captures the forward, loss, and backward step as one graph that can be transformed, precompiled once, and reused across ranks. Its graph passes manage activation memory and overlap FSDP or expert-parallel communication with computation.

GraphPP partitions this graph into pipeline stages and supports schedules including DualPipeV.

Key PRs:

• Whole-step tracing (#2561)
• Rank-agnostic precompilation (#2713)
• Activation-memory framework (#3118)
• FSDP communication overlap (#3156)
• EP communication overlap (#3328)
• CUDA Graph execution (#2868)
• GraphPP stage partitioning (#2726)
• Backward splitting (#2727)
• FSDP collective splitting (#3088)
• GraphPP runner (#3089)
• DualPipeV validation (#3090)

Precision, memory, and training efficiency

BF16 optimizer states, MXFP8, Float8, and NVFP4 reduce memory use and enable lower-precision training. Composable activation-checkpointing policies reduce activation memory, while Trainer CUDA graphs reduce CPU launch overhead.

Key PRs:

• BF16 optimizer states (#2732)
• MXFP8 conversion (#2520)
• NVFP4 (#3914)
• Checkpointing policies (#3674)
• CUDA graphs (#3559)

Supervised fine-tuning

ChatDataset and a new SFT recipe support training on chat-formatted conversations.

Key PR: ChatDataset and SFT (#2556)

TitanRL

Experimental TitanRL workflows add asynchronous and multi-turn rollouts, continuous batching, multi-generator routing, overlapped weight synchronization, checkpoint and resume support, and Search-R1 and DAPO recipes.

TitanRL requires additional dependencies. Follow the TitanRL quick-start guide before running these workflows.

Key PRs:

• Async RL loop (#3642)
• Continuous batching and multi-turn rollouts (#3593)
• Overlapped weight sync (#3810)
• Checkpoint and resume (#3751)
• Search-R1 (#3602)
• DAPO (#3951)

Compatibility

Installation

uv pip install torchtitan==0.3.0

Dependency

PyTorch Version: torch-2.12.0.dev20260220+cu126
TorchAO Version: torchao-0.17.0.dev20260220+cu126

What's changed

🚀 Features

• [CP] Refactor Context Parallel to use new PyTorch CP APIs (#2144) by @fegin
• [CP] Enable FlexCP for llama3 (#2145) by @fegin
• [Compiler Toolkit] Add option for full inductor (#2150) by @aditvenk
• Add docs to explain COMM_MODE (#2162) by @fegin
• Disable dynamo LRU cache when AC is enabled (#2204) by @soulitzer
• feat(gpt-oss): add YaRN RoPE extensions with mscale for extended context (#2216) by @eous
• [ROCm] Support mxfp8 on gfx950 (#2222) by @RuibinCheung
• Enable memory snapshot for generic devices (#2228) by @frost-intel
• GQA without kv repeats (#2259) by @francesco-bertolotti
• [rl] GQA attention enablement in torchtitan vllm wrapper (#2299) by @wwwjn
• Add peak flops for NVIDIA H20 GPUs (#2307) by @DamonFool
• Add missing job_config.maybe_log() calls (#2308) by @EquationWalker
• [DeepEP] Implement shared_experts overlap with deepep.combine() (#2310) by @vivekgoe
• Separate out training for fault tolerance (#2311) by @tushar00jain
• Torchtitan changes to integrate into Verl (#2333) by @acisseJZhong
• Maintain same LR schedule for early stop debug runs (#2340) by @acisseJZhong
• [Compiler Toolkit] Separate process groups for FSDP AG/RS comm overlap (#2368) by @yiming0416
• [AC] Set preserve_rng_state=True as default for activation checkpointing (#2380) by @soulitzer

🧠 Model

• Add attention scaling to varlen for Qwen3 (#2178) by @liangel-02
• Make get TP mesh optional in Llama4 parallelize (#2185) by @danielvegamyhre
• [GPT-OSS] Graduate from experiments to main (#2203) by @shuhuayu
• [autoparallel] Update local_map_deepseek_v3 device mesh usage (#2231) by @xmfan
• [varlen_attn] Change is_causal to window_size (#2267) by @liangel-02
• Remove _ScaledPartial placement (#2337) by @Aidyn-A
• Added custom trunc_normal (#2342) by @francesco-bertolotti
• Removed weight initialization from model __init__ (#2361) by @francesco-bertolotti

🐛 Bug Fixes

• [docs] Fix missing --model.flavor flags in compiler_toolkit README (#2201) by @BryanBradfo
• Fix loss computation by handling valid token imbalance in train loop (#2206) by @wwwjn
• [MoE] Fix experts DTensor metadata bug for DCP (#2227) by @shuhuayu
• Fix sdpa-varlen attention mismatch in Qwen3 (#2229) by @francesco-bertolotti
• Weight tying fix for Qwen3 (#2253) by @francesco-bertolotti
• Fix grad norm clipping for AutoP and DSv3 model init (#2270) by @sanketpurandare
• [MoE] DeepEP refactor and fix memory leak during training and inference (#2296) by @shuhuayu
• [docs] Fix type mismatch in model layers comments (#2306) by @DamonFool
• Fix FLUX attention by exposing is_causal in SDPA (#2309) by @wwwjn
• Fixing global_max_loss computation (#2314) by @Shagun-G
• Fix the CI loss issue (#2315) by @fegin
• Fix gpt-oss implementation (MoE router gate bias + top-k renorm) (#2319) by @linyuhongg
• [SimpleFSDP] Fix HSDP placement mismatch in _distribute_dtensor (#2329) by @SongyuanZhao
• [Bugfix] Fix bitwise determinism after vLLM SiluAndMul change (#2358) by @Lucaskabela
• [Bugfix] Fix simple_rl_multiprocess.py to be runnable with recent vLLM version (#2359) by @Lucaskabela
• Fixing extra averaging performed in validation error (#2366) by @Shagun-G
• Bug fix: don't swallow OutOfMemoryError when enable_memory_snapshot=True (#2374) by @weifengpy
• Fix: restrict completion logging to rank 0 (#2383) by @fatih-uzlmz

🧪 Experiments / CI / Infra

• [Experimental][rl][vllm compat] Update simple_rl example to work with vLLM nightly (#2219) by @Lucaskabela
• [Experimental][rl][unified] Update infer.py example to work with vLLM nightly (#2226) by @Lucaskabela
• Add ROCm support for H100 tests (#2202) by @akashveramd
• Add ROCm CI support for simple FSDP experiments test (#2220) by @akashveramd
• Add test for DSv3 with flexattn + fsdp + ep + pp + sac op (#2234) by @shuhuayu
• Add ROCm CI support for Auto Parallel & Compiler Toolkit experiments (#2248) by @akashveramd
• Add ROCm CI support for Transformers Modeling Backend & VLM experiments (#2276) by @akashveramd
• Update CPU unit test to use linux_job_v2 (#2287) by @joecummings
• [BE week] Disable CPU wheel builds in nightly CI (#2289) by @joecummings
• [BE][NFC] Add integration test for simplefsdp + CP deepseek_v3 (#2301) by @aditvenk
• Fixed autoparallel integration tests on ROCm (#2321) by @wenchenvincent
• [rl][ez] Squash landing import and git fixes (#2331) by @zhxchen17
• [ci] Add DSv3 SimpleFSDP auto_bucketing to H100 CI jobs (#2347) by @IvanKobzarev
• Bump tj-actions/changed-files from 47.0.1 to 47.0.2 (#2367) by @dependabot
• [CI] Disable NVLS (#2372) by @fegin
• Bump tj-actions/changed-files from 47.0.2 to 47.0.4 (#2390) by @dependabot
• [rl] Install vllm from pre-built wheels (#2397) by @wwwjn
• [rl/unified] Update default model-ckpt-path in infer.py to the one from README (#2405) by @daniellepintz

🔧 Typing / Lint Cleanup

• [lint] Ignore all existing pyrefly errors (#2240) by @xmfan
• [Typing] Fix CI Typing Issues (#2245) by @fegin
• [Typing] Improve ModelProtocol typing (#2246) by @fegin
• [Typing] Remove deprecated enable_symm_mem_for_group (#2260) by @fegin
• [Typing] Remove unused pyrefly ignore (#2280) by @fegin
• [Typing] Fix pyrefly-ignore in train.py (#2282) by @fegin
• [Typing] Fix pyrefly ignores in checkpoint.py (#2283) by @fegin
• [Typing] Fix the ignores in activation_checkpoint.py (#2284) by @fegin
• [Typing] Fix the ignores in tokenizer.py (#2285) by @fegin
• [Typing] Fix the ignores in validate.py (#2286) by @fegin
• [Typing] Fix some pyrefly ignores in optimizer.py (#2294) by @fegin
• [Typing] Improve typing for some distributed modules (#2295) by @fegin
• [Typing] Fix the pyrefly ignores in llama3 model.py (#2302) by @fegin
• [Typing] Fix pyrefly ignores in llama4 model.py (#2303) by @fegin
• [Typing] Fix pyrefly ignores in qwen3 model.py (#2304) by @fegin
• [Typing] Fix pyrefly ignores in deepseek model.py (#2305) by @fegin

Dependency

pytorch verison: torch-2.11.0.dev20251226+cu126
torchao version: torchao-0.16.0.dev20251226+cu126

What's Changed

Features

• Use new DeviceMesh unflatten to rewrite parallel_dims by @fegin in #1660
• Re:Run Torchtitan ROCm workflow on cron schedule & push to Main branch only by @akashveramd in #2018
• adding variable length attention to llama3 8b by @liangel-02 in #2000
• [Local Tensor] Replace dry_run.py with fake mode implementation by @fegin in #2057

Model

• Enable PP and EP overlap for MoE by @H-Huang in #1721
• Integrate DeepEP to torchtitan by @elfiegg in #2107
• [MoE] Add node limited routing support by @shuhuayu in #2111
• Add Context Parallelism to Flux model training by @limou102 in #1851
• gpt-oss model enablement by @wwwjn in #1754
• [GPT-OSS] Add HF state dict adapter to support loading from HF checkpoints by @shuhuayu in #2021

Bug Fix

• [FLOPs] Fix attention FLOPs estimate by @shuhuayu in #1923
• Fix apply_compile called multiple times in PP initialization by @xmfan in #2135
• Fix qwen3 attention scaling calculation by @wwwjn in #2173

Experiments

• Exploring toolkit-style use of the compiler stack @SherlockNoMad @yiming0416 : [Compiler Toolkit] JointGraph-based Training Prototype for llama3 by @SherlockNoMad in #1794
• Bit-wise identity RL between torchtitan Trainer and vLLM sampler: Add deterministic RL training experiment with vLLM by @bwasti in #1975
• Train models fromtransformers with torchtitan: 3outeille/transformers backend (Dense model only) by @3outeille in #2048
• Auto Parallel Examples @wconstab @xmfan: Autoparallel as an experiment in main by @xmfan in #2054
• Unified Model definition in RL loop @wwwjn @acisseJZhong @zhxchen17 : Run vLLM inference using torchtitan model definition (single GPU) by @wwwjn in #2119

Full Changelog: v0.2.0...v0.2.1

Dependency

pytorch verison: torch-2.10.0.dev20251019+cu126
torchao version: torchao-0.15.0.dev20251015+cu126

Full Changelog: v0.1.0...v0.2.0

This is the first pre-release of torchtitan, following the release practice outlined in #688.

torch version: torch-2.8.0.dev20250617+cu126
torchao version: torchao-0.12.0.dev20250617+cu126