torchtitan

pytorchBSD-3-Clause活跃维护

PyTorch 官方的生成式 AI 模型训练平台,提供原生分布式训练组件和可复用的并行策略,帮你在多卡集群上高效训练大模型。

A PyTorch native platform for training generative AI models

v0.4.0发布于 2026-10-28 18:05 UTC5,889star +5 / 7 天1,077fork同步于 2026-11-10 11:35 UTC

怎么上手

由 README 提炼,执行前请确认
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
硬件要求按 README 说明

README 未给出具体显存要求;快速开始默认单卡可跑 debug 模型,真实模型示例按 8 GPU 配置,CUDA 与 ROCm 构建均支持,AMD 另有官方优化 fork

需要 CUDA支持 ROCm / AMD多卡 / 张量并行
下载:源码 zip·GitHub Releases(安装包)本站不镜像安装包

项目说明

由 README 自动翻译提炼

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 等模型

容易踩的坑

需搭配较新的 PyTorch nightly,稳定版 torchtitan 会锁定其验证过的 torch/torchao nightly 版本;MXFP8/NVFP4 所需的 torchao 不在 requirements.txt 里,要按对应 torch 构建单独安装;跑真实模型需先用 scripts/download_hf_assets.py 下载 tokenizer,Llama 权重还需在 Hugging Face 申请访问权限

内容由模型根据 README 生成,可能不完整或不准确,以官方文档为准。看原文 README

最近发版

5 个版本v0.4.0
  1. 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_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
  2. 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:

    • 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:

    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:

  3. v0.2.2pre-release2026-02-20 · 2026-02-20 22:47 UTC
    查看更新说明

    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
  4. v0.2.1pre-release2025-12-26 · 2025-12-26 23:29 UTC
    查看更新说明

    Dependency

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

    What's Changed

    Features

    Model

    Bug Fix

    Experiments

    Full Changelog: https://github.com/pytorch/torchtitan/compare/v0.2.0...v0.2.1

  5. v0.2.0pre-release2025-10-18 · 2025-10-18 04:26 UTC
    查看更新说明

    Dependency

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

    Full Changelog: https://github.com/pytorch/torchtitan/compare/v0.1.0...v0.2.0

用户点评

登录用户的真实使用记录,可按版本和硬件筛选
共 0 条

    还没有人点评。第一个说出真实体验的人最有价值。

    README 原文

    未翻译

    A PyTorch-native platform for pretraining and post-training generative AI models at scale

    arXiv ICLR devlogs license pip conda

    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 DeviceMesh and 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] torchtitan is 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 torchtitan for AMD GPUs.
    • [2025/10] We released torchtitan v0.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

    ModelSizesModalityVerified recipesRL
    Llama 31B – 405BTextllama3_8b, llama3_70b✅
    Qwen30.6B – 32B dense; 30B-A3B, 235B-A22B MoETextqwen3_14b✅
    Qwen3.5 / 3.6 / 3.80.8B – 397B-A17B; 2.4T-A95B (3.8)Text, vision✅ text
    DeepSeek V316B, 236B, 671BTextdeepseek_v3_671b (+ Dist-MoE, NVFP4)
    DeepSeek V4Flash, ProTextSoon
    gpt-oss20B, 120BText✅
    Kimi K2.7Moonlight-16B-A3B, Kimi-VL-A3B, K2.5Text, vision
    Kimi K3K3Text, visionSoon
    Muse Glimmer30BText, visionmuse_glimmer_30b✅
    Muse-Spark-Open816B-A42BText, vision, audiomuse_spark_816b_a42b_dist_moe_bf16_spmd_256✅ text
    FLUXdev, schnellText-to-imageflux_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.pipelining schedules
    • 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

    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

    Checkpointing

    Debugging and observability

    Extensibility

    We report performance on up to 512 GPUs and verify loss convergence across parallelisms and techniques.

    Test status

    HardwareIntegration testsUnit tests
    CPUCPU Unit Test
    NVIDIA GPUIntegration Tests H100 Tests B200 TestsGPU Unit Tests
    AMD GPU (ROCm)Integration Tests (ROCm)
    RL (NVIDIA GPU)RL Integration Tests RL B200 TestsRL GPU Unit Tests

    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 torchao nightly CI pins, from the index matching your torch build. It is not in requirements.txt so it cannot be resolved against a different torch:
      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.sh installs 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 torchtitan from 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

    TopicWhere
    Configuring runs and writing recipestorchtitan/config
    Adding a modeltorchtitan/models
    Parallelisms and distributed techniquestorchtitan/distributed
    Datasets and data loadingtorchtitan/components/data
    Checkpointing and HF conversiontorchtitan/components/checkpointer
    Reinforcement learningtorchtitan/rl
    Debugging, reproducibility, and validating numericsdocs/debugging.md
    Validation and evaluationdocs/evaluation.md
    Benchmarksdocs/benchmarks
    Tests and CItests

    Contributing

    We look forward to your contributions!

    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.