需搭配较新的 PyTorch nightly,稳定版 torchtitan 会锁定其验证过的 torch/torchao nightly 版本;MXFP8/NVFP4 所需的 torchao 不在 requirements.txt 里,要按对应 torch 构建单独安装;跑真实模型需先用 scripts/download_hf_assets.py 下载 tokenizer,Llama 权重还需在 Hugging Face 申请访问权限
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.shREADME 未给出具体显存要求;快速开始默认单卡可跑 debug 模型,真实模型示例按 8 GPU 配置,CUDA 与 ROCm 构建均支持,AMD 另有官方优化 fork
PyTorch 原生的大模型训练平台,从预训练到强化学习一套代码搞定
PyTorch 官方的生成式 AI 模型训练平台,提供原生分布式训练组件和可复用的并行策略,帮你在多卡集群上高效训练大模型。
核心能力
- 用 DeviceMesh 和 DTensor 叠加 FSDP2、TP、PP、CP、EP 组成 N 维并行
- 每次训练都是一个可读可改的 Python recipe 函数,不用堆命令行参数
- TitanRL 在训练与 vLLM 生成间复用同一模型定义,支持 batch-invariant 模式
- 内置 MXFP8/NVFP4 低精度训练与 Dist-MoE、CUDA graphs 等前沿技术
- 覆盖 Llama、Qwen、DeepSeek、Kimi、gpt-oss、FLUX 等模型
容易踩的坑
- v0.4.0最新正式版2026-10-28 · 2026-10-28 18:05 UTC
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_regionsselects functions such as SwiGLU, RoPE, and gated RMSNorm for independent compilation and kernel fusion withtorch.compile. (#4894, #4931, #4985, #5026) - Faster router and LM-head layers:
HiMidLoLinearuses 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_rematcontrols 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
BlockShardpartitions and BF16 redistribution
- v0.3.02026-09-03 · 2026-09-03 03:51 UTC
查看更新说明
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:
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_typesis 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_dtensorbackend, 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](https://github.com/pytorch/torcht
- v0.2.2pre-release2026-02-20 · 2026-02-20 22:47 UTC
查看更新说明
Dependency
PyTorch Version:
torch-2.12.0.dev20260220+cu126TorchAO Version:torchao-0.17.0.dev20260220+cu126What'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=Trueas 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_v3device mesh usage (#2231) by @xmfan - [varlen_attn] Change
is_causaltowindow_size(#2267) by @liangel-02 - Remove
_ScaledPartialplacement (#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.flavorflags 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_causalin SDPA (#2309) by @wwwjn - Fixing
global_max_losscomputation (#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
SiluAndMulchange (#2358) by @Lucaskabela - [Bugfix] Fix
simple_rl_multiprocess.pyto 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
OutOfMemoryErrorwhenenable_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.pyexample 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
- v0.2.1pre-release2025-12-26 · 2025-12-26 23:29 UTC
查看更新说明
Dependency
pytorch verison:
torch-2.11.0.dev20251226+cu126torchao version:torchao-0.16.0.dev20251226+cu126What's Changed
Features
- Use new DeviceMesh unflatten to rewrite parallel_dims by @fegin in https://github.com/pytorch/torchtitan/pull/1660
- Re:Run Torchtitan ROCm workflow on cron schedule & push to Main branch only by @akashveramd in https://github.com/pytorch/torchtitan/pull/2018
- adding variable length attention to llama3 8b by @liangel-02 in https://github.com/pytorch/torchtitan/pull/2000
- [Local Tensor] Replace dry_run.py with fake mode implementation by @fegin in https://github.com/pytorch/torchtitan/pull/2057
Model
- Enable PP and EP overlap for MoE by @H-Huang in https://github.com/pytorch/torchtitan/pull/1721
- Integrate DeepEP to torchtitan by @elfiegg in https://github.com/pytorch/torchtitan/pull/2107
- [MoE] Add node limited routing support by @shuhuayu in https://github.com/pytorch/torchtitan/pull/2111
- Add Context Parallelism to Flux model training by @limou102 in https://github.com/pytorch/torchtitan/pull/1851
- gpt-oss model enablement by @wwwjn in https://github.com/pytorch/torchtitan/pull/1754
- [GPT-OSS] Add HF state dict adapter to support loading from HF checkpoints by @shuhuayu in https://github.com/pytorch/torchtitan/pull/2021
Bug Fix
- [FLOPs] Fix attention FLOPs estimate by @shuhuayu in https://github.com/pytorch/torchtitan/pull/1923
- Fix apply_compile called multiple times in PP initialization by @xmfan in https://github.com/pytorch/torchtitan/pull/2135
- Fix qwen3 attention scaling calculation by @wwwjn in https://github.com/pytorch/torchtitan/pull/2173
Experiments
- Exploring toolkit-style use of the compiler stack @SherlockNoMad @yiming0416 : [Compiler Toolkit] JointGraph-based Training Prototype for llama3 by @SherlockNoMad in https://github.com/pytorch/torchtitan/pull/1794
- Bit-wise identity RL between torchtitan Trainer and vLLM sampler: Add deterministic RL training experiment with vLLM by @bwasti in https://github.com/pytorch/torchtitan/pull/1975
- Train models from
transformerswith torchtitan: 3outeille/transformers backend (Dense model only) by @3outeille in https://github.com/pytorch/torchtitan/pull/2048 - Auto Parallel Examples @wconstab @xmfan: Autoparallel as an experiment in main by @xmfan in https://github.com/pytorch/torchtitan/pull/2054
- Unified Model definition in RL loop @wwwjn @acisseJZhong @zhxchen17 : Run vLLM inference using torchtitan model definition (single GPU) by @wwwjn in https://github.com/pytorch/torchtitan/pull/2119
Full Changelog: https://github.com/pytorch/torchtitan/compare/v0.2.0...v0.2.1
- v0.2.0pre-release2025-10-18 · 2025-10-18 04:26 UTC
查看更新说明
Dependency
pytorch verison:
torch-2.10.0.dev20251019+cu126torchao version:torchao-0.15.0.dev20251015+cu126Full Changelog: https://github.com/pytorch/torchtitan/compare/v0.1.0...v0.2.0
用户点评
登录用户的真实使用记录,可按版本和硬件筛选还没有人点评。第一个说出真实体验的人最有价值。
A PyTorch-native platform for pretraining and post-training generative AI models at scale
Quick start · Models · Features · RL · Docs · Contributing
torchtitan is a clean-room, PyTorch-native implementation of large-scale
training. One readable codebase takes a model from pretraining through
supervised fine-tuning to reinforcement learning, composing FSDP, tensor,
pipeline, context, and expert parallelism without rewriting the model.
- Composable parallelism. Models declare how they shard; FSDP2, TP, PP, CP,
and EP stack into N-D layouts through
DeviceMeshand DTensor. - Built for research. Minimal, hackable code; every run is a Python recipe you can read, copy, and modify.
- Pretraining to RL in one stack. TitanRL runs the same model definitions and kernels in the trainer and in vLLM generation, with a bitwise-reproducible batch-invariant mode.
- Frontier models and techniques. DeepSeek, Qwen, Kimi, gpt-oss, Llama, and FLUX, with MXFP8/NVFP4 low precision, Dist-MoE, CUDA graphs, and more.
[!NOTE]
torchtitanis under active development. To use the latest features, use a recent PyTorch nightly.
Latest news
- [2026/08] TitanRL is a hackable RL stack for scaling and debugging. It reuses TorchTitan model definitions and kernels across training and vLLM generation and supports batch-invariant mode.
- [2025/11] AMD released an optimized fork of
torchtitanfor AMD GPUs. - [2025/10] We released
torchtitanv0.2.0. - [2025/10] SkyPilot now supports
torchtitan! See the tutorial.
- [2025/07] We published instructions on how to add a model to
torchtitan. - [2025/04] Our paper was accepted by ICLR 2025.
- [2024/12] GPU MODE lecture on torchtitan.
- [2024/07] Presentation at PyTorch Conference 2024.
Quick start
1. Install with a PyTorch nightly (Python 3.11+). Replace cu132 with your CUDA or ROCm build.
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
2. Train a debug model. It uses the tokenizer bundled with the tests, so no download is needed.
NGPU=1 ./run_train.sh # Llama 3 debug model, 10 steps
3. Train a real model. Download the tokenizer (Llama weights require access on Hugging Face), then select a recipe by module and function name:
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...
MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8b ./run_train.sh # 8 GPUs
4. Make it yours. A recipe is a Python function returning a complete config, so changing a run means writing a function, not passing flags:
# 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
MODULE=my_recipes CONFIG=llama3_8b_tp2 ./run_train.sh
See the configuration guide for the full model,
and torchtitan_recipes for the verified recipes.
Supported models
| Model | Sizes | Modality | Verified recipes | RL |
|---|---|---|---|---|
| Llama 3 | 1B – 405B | Text | llama3_8b, llama3_70b | ✅ |
| Qwen3 | 0.6B – 32B dense; 30B-A3B, 235B-A22B MoE | Text | qwen3_14b | ✅ |
| Qwen3.5 / 3.6 / 3.8 | 0.8B – 397B-A17B; 2.4T-A95B (3.8) | Text, vision | ✅ text | |
| DeepSeek V3 | 16B, 236B, 671B | Text | deepseek_v3_671b (+ Dist-MoE, NVFP4) | |
| DeepSeek V4 | Flash, Pro | Text | Soon | |
| gpt-oss | 20B, 120B | Text | ✅ | |
| Kimi K2.7 | Moonlight-16B-A3B, Kimi-VL-A3B, K2.5 | Text, vision | ||
| Kimi K3 | K3 | Text, vision | Soon | |
| Muse Glimmer | 30B | Text, vision | muse_glimmer_30b | ✅ |
| Muse-Spark-Open | 816B-A42B | Text, vision, audio | muse_spark_816b_a42b_dist_moe_bf16_spmd_256 | ✅ text |
| FLUX | dev, schnell | Text-to-image | flux_dev, flux_schnell |
Every model has debug configurations exercised in CI. Verified recipes have been run in training, convergence, or performance workflows. To add a model, see the model guide.
Features
Parallelism
- FSDP2, HSDP, and DDP
- Tensor parallel, including async TP
- Pipeline parallel: 1F1B, Interleaved 1F1B, GPipe, and other
torch.distributed.pipeliningschedules - Context parallel: all-gather KV and Ulysses
- Expert parallel with DeepEP / HybridEP dispatch
- Dist-MoE: fused dispatch, expert compute, and combine on Blackwell
Memory and performance
- Activation checkpointing: full, selective, and regional
- Regional
torch.compileand CUDA graphs - MXFP8 and NVFP4 training
- FlashAttention 3/4, FlexAttention, and varlen attention
- BF16 optimizer states and FSDP CPU offload
Training
- Grain data pipeline with packing, mixing, and checkpointable state
- Pretraining, SFT on chat datasets, and LoRA
- Muon (DistMuon), EMA, multi-token prediction
- Warmup-stable-decay LR schedule and validation
Post-training with TitanRL
- One model definition shared by the trainer and vLLM
- Bitwise trainer/generator parity in batch-invariant mode
- Async rollouts orchestrated with Monarch; trainer-to-generator model state dict transfer over TorchStore
- GRPO and DAPO losses; DAPO Math, Search-R1, and Verifiers examples
Checkpointing
- Distributed checkpointing, including async saves
- Load and save Hugging Face safetensors directly
- Conversion scripts between HF and DCP
Debugging and observability
- Metrics: loss, memory, throughput, TFLOPs, and MFU in TensorBoard or W&B
- Profiling and memory snapshots
- Flight Recorder and other debugging tools
- Fake-backend dry runs of large layouts on one GPU
- Deterministic mode and loss comparison across commits
- SDC replay and structured logging
Extensibility
- Python recipes selected with
--moduleand--config - Component overrides for custom kernels and modules
- Experiments: GraphTrainer, TorchFT, Transformers modeling backend
We report performance on up to 512 GPUs and verify loss convergence across parallelisms and techniques.
Test status
| Hardware | Integration tests | Unit tests |
|---|---|---|
| CPU | ||
| NVIDIA GPU | ||
| AMD GPU (ROCm) | ||
| RL (NVIDIA GPU) |
Installation
The quick start installs from source, which is the recommended way to get the latest features. Other options:
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu132 --force-reinstall
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu132
Replace cu132 with another CUDA version or a ROCm build (e.g. rocm10.0).
pip install torchtitan
# or
conda install conda-forge::torchtitan
Each stable release pins the nightly versions of torch and torchao it was
validated against; see RELEASE.md.
- Low precision (MXFP8, NVFP4): install the
torchaonightly CI pins, from the index matching yourtorchbuild. It is not inrequirements.txtso it cannot be resolved against a differenttorch:USE_CPP=0 pip install --pre -r .github/requirements-torchao.txt --index-url https://download.pytorch.org/whl/nightly/cu132 - RL:
scripts/install_rl_env.shinstalls the vLLM, Monarch, and TorchStore nightlies and the FlashAttention build for your GPU. See the RL quick start. - Extras for vision, FLUX, RL, and development are also available as
pip install -e ".[vlm,flux,rl,dev]". - Importing from elsewhere: the source tree runs as is; to import
torchtitanfrom another directory, install it in editable mode without re-resolving dependencies:pip install -e . --no-deps.
Submit scripts/multinode_trainer.slurm with
sbatch from the repository root. Set the node count in both the #SBATCH
header and the torchrun --nnodes argument:
#SBATCH --ntasks=2
#SBATCH --nodes=2
srun torchrun --nnodes 2 ...
If a node does not have 8 GPUs, also adjust --nproc_per_node and
#SBATCH --gpus-per-task.
Documentation
| Topic | Where |
|---|---|
| Configuring runs and writing recipes | torchtitan/config |
| Adding a model | torchtitan/models |
| Parallelisms and distributed techniques | torchtitan/distributed |
| Datasets and data loading | torchtitan/components/data |
| Checkpointing and HF conversion | torchtitan/components/checkpointer |
| Reinforcement learning | torchtitan/rl |
| Debugging, reproducibility, and validating numerics | docs/debugging.md |
| Validation and evaluation | docs/evaluation.md |
| Benchmarks | docs/benchmarks |
| Tests and CI | tests |
Contributing
We look forward to your contributions!
- New ideas should start in the
experimentsfolder; follow the experiments guidelines. - For fixes and contributions to core, follow the contributing guidelines.
Citation
Our paper gives a detailed look at the parallelisms and optimizations in
torchtitan, with advice on when to use each technique:
TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training.
@inproceedings{
liang2025torchtitan,
title={TorchTitan: One-stop PyTorch native solution for production ready {LLM} pretraining},
author={Wanchao Liang and Tianyu Liu and Less Wright and Will Constable and Andrew Gu and Chien-Chin Huang and Iris Zhang and Wei Feng and Howard Huang and Junjie Wang and Sanket Purandare and Gokul Nadathur and Stratos Idreos},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=SFN6Wm7YBI}
}
License
Source code is made available under a BSD 3 license. You may have other legal obligations that govern your use of other content linked in this repository, such as the license or terms of service for third-party data and models.