torchtitan 目前正处于大规模开发中。为了使用 torchtitan 的最新功能，我们建议您使用最新的 PyTorch 夜间版本。

最新消息

• [2025年11月] AMD 发布了针对 AMD 显卡的 torchtitan 优化分支。
• [2025年10月] 我们发布了 torchtitan v0.2.0。
• [2025年10月] SkyPilot 现已支持 torchtitan！请参阅教程 这里。
• [2025年7月] 我们发布了关于如何将模型添加到 torchtitan 的 说明。
• [2025年4月] 我们的论文已被 ICLR 2025 接受。
• [2024年12月] GPU MODE 讲座 讲解 torchtitan。
• [2024年7月] 在 PyTorch 2024 大会上进行了 演讲。

概述

torchtitan 是一个原生基于 PyTorch 的平台，专为生成式 AI 模型的 快速实验和大规模训练 而设计。作为 PyTorch 原生扩展技术的一个极简、干净的实现，torchtitan 为开发者提供了一个灵活的基础，便于在此之上进行构建。借助 torchtitan 的 扩展点，用户可以轻松创建满足特定需求的自定义扩展。

我们的使命是通过赋能研究人员和开发者探索新的模型架构与基础设施技术，加速生成式 AI 领域的创新。

构建 torchtitan 时的指导原则：

• 设计简洁易懂，易于使用且可针对不同训练目的进行扩展。
• 应用多维并行时对模型代码的改动最小化。
• 偏向于保持代码库的整洁与精简，同时提供基础的可重用或可替换组件。

torchtitan 已通过支持预训练各种规模的 Llama 3.1 LLM，展示了 PyTorch 最新的分布式训练特性。

参与贡献

我们期待您的贡献！

• 为了加速对 torchtitan 的贡献及围绕它的创新，我们设立了 experiments 文件夹。新想法应从那里开始。如需贡献，请遵循 experiments 指南。
• 对核心部分的修复与贡献，请遵循这些 指南。

Llama 3.1 训练

可用的关键特性

1. 多维可组合并行
   • 带有参数级分片的 FSDP2
   • 张量并行（包括 异步 TP）
   • 流水线并行
   • 上下文并行
2. Meta device 初始化
3. 按操作选择性及全激活检查点保存
4. 分布式检查点（包括异步检查点）
   • 互操作检查点，可以直接加载到 torchtune 中用于微调
5. torch.compile 支持
6. Float8 支持（使用方法）
7. Blackwell GPU 上的 MXFP8 训练，适用于密集型和 MoE 模型。
8. DDP 和 HSDP
9. TorchFT 集成
10. 可检查点的数据加载，预配置了 C4 数据集（1.44亿条记录），并支持 自定义数据集
11. 梯度累积，可通过在命令行中添加额外的 --training.global_batch_size 参数来启用
12. 灵活的学习率调度器（预热-稳定-衰减）
13. 损失、GPU 内存、吞吐量（tokens/sec）、TFLOPs 和 MFU 会通过 TensorBoard 或 Weights & Biases 进行显示和记录
14. 调试工具，包括 CPU/GPU 性能分析、内存分析、Flight Recorder 等
15. 所有选项均可通过 Python 配置注册表 轻松配置，只需使用 --module 和 --config 命令行参数
16. 辅助脚本 用于：
    • 从 Hugging Face 下载分词器
    • 将原始 Llama 3 检查点转换为预期的 DCP 格式
    • 在不实例化模型的情况下估算 FSDP/HSDP 的内存占用
    • 使用张量并行运行分布式推理

我们报告了在最多 512 张 GPU 上的 性能，并验证了各种技术的 损失收敛 正确性。

深入代码

你可能想了解模型是如何定义的，或者并行化技术是如何应用的。为了更好地理解，可以先查看以下文件：

• torchtitan/train.py - 主训练循环及高层次的设置代码
• torchtitan/models/llama3/model.py - Llama 3.1 模型的定义
• torchtitan/models/llama3/parallelize.py - 用于将数据并行、张量并行、激活检查点以及 torch.compile 应用到模型上的辅助工具
• torchtitan/distributed/pipeline_parallel.py - 用于将流水线并行应用到模型上的辅助工具
• torchtitan/components/checkpoint.py - 用于保存和加载分布式检查点的工具
• torchtitan/components/quantization/float8.py - 用于应用 Float8 技术的工具

安装

你可以直接运行源代码，也可以从 nightly 构建版本或稳定版中安装 torchtitan。

从源码安装

此方法需要 PyTorch 的 nightly 构建版本，或者从 源码 构建的最新 PyTorch 版本。

git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt
pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu

注意： 使用 PyTorch 的 nightly 版本时，需要 torchdata 的 nightly 构建版本。请按照上述方式从 nightly 索引中安装。

Nightly 构建版本

此方法需要 PyTorch 的 nightly 构建版本。你可以将 cu128 替换为其他版本的 CUDA 或 AMD GPU（例如 rocm6.3）。

pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 --force-reinstall
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu128

安装稳定版

你可以通过 pip 或 conda 安装 torchtitan 的最新 稳定版。

pip install torchtitan

conda install conda-forge::torchtitan

请注意，每个稳定版都会锁定 torch 和 torchao 的 nightly 版本。更多详情请参阅 release.md 文件。

下载分词器

torchtitan 目前开箱即用地支持训练 Llama 3.1（8B、70B、405B）。要开始训练这些模型，我们需要下载分词器。请按照官方 meta-llama 仓库中的说明操作，以确保你有权访问 Llama 模型权重。

确认权限后，你可以运行以下命令将 Llama 3.1 分词器下载到本地机器上。

# 从 https://huggingface.co/settings/tokens 获取你的 HF 令牌

# Llama 3.1 分词器
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

开始训练

在 8 张 GPU 上本地训练 Llama 3 8B 模型

MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

多节点训练

对于 ParallelCluster/Slurm 类型的配置，你可以使用 multinode_trainer.slurm 文件提交 sbatch 作业。

首先调整节点和 GPU 的数量：

#SBATCH --ntasks=2
#SBATCH --nodes=2

然后启动一个运行，其中 nnodes 是你的总节点数，与上面的 sbatch 节点数一致。

srun torchrun --nnodes 2

如果你每节点的 GPU 数量不是 8 张，请相应调整 torchrun 命令中的 --nproc_per_node 以及 SBATCH 命令部分的 #SBATCH --gpus-per-task。

引用

我们详细介绍了 torchtitan 中可用的并行化和优化技术，并总结了何时使用各种技术的建议。

TorchTitan：面向生产级 LLM 预训练的一站式 PyTorch 原生解决方案

@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}
}

许可证

源代码根据 BSD 3 许可证 提供，但你可能还需遵守其他法律义务，这些义务适用于本仓库中链接的其他内容，例如第三方数据和模型的许可证或服务条款。

TorchTitan 快速上手指南

TorchTitan 是一个原生的 PyTorch 平台，专为生成式 AI 模型的快速实验和大规模训练而设计。它提供了最小化且清晰的代码实现，支持多维并行策略（如 FSDP2、张量并行、流水线并行等），是预训练 Llama 3.1 等大模型的理想选择。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• GPU: NVIDIA GPU (支持 CUDA) 或 AMD GPU (需使用特定 fork 版本)
• Python: 3.8 或更高版本

前置依赖

• PyTorch:
  • 若要使用最新特性，强烈建议安装 PyTorch Nightly 版本。
  • 若使用稳定版，请注意部分新特性可能不可用。
• 其他依赖: torchdata (Nightly 版本需单独安装)。

注意: 国内用户若访问 PyT官方源较慢，可尝试配置清华或阿里镜像源加速 pip 下载，但 Nightly 版本通常仅官方源提供最新构建。

安装步骤

你可以选择从源码安装、安装 Nightly 版本或稳定版本。

方式一：从源码安装（推荐用于开发和新特性体验）

此方法需要预先安装 PyTorch Nightly。

# 1. 克隆仓库
git clone https://github.com/pytorch/torchtitan
cd torchtitan

# 2. 安装基础依赖
pip install -r requirements.txt

# 3. 安装 PyTorch Nightly 和 torchdata (以 CUDA 12.8 为例，可根据实际环境调整 cu12x)
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 --force-reinstall
pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu

方式二：安装稳定发布版

适合生产环境或不需要最新实验特性的场景。

使用 pip:

pip install torchtitan

使用 conda:

conda install conda-forge::torchtitan

提示: 稳定版本会锁定特定的 torch 和 torchao 夜间版版本，详情请查阅项目文档中的 release.md。

获取模型分词器 (Tokenizer)

TorchTitan 原生支持 Llama 3.1 (8B, 70B, 405B) 训练。开始前需下载分词器。 请确保你已在 Hugging Face 获得 meta-llama 的访问权限并获取 Token。

# 将 ... 替换为你的 Hugging Face Token
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

基本使用

单机多卡训练示例

以下命令演示如何在本地 8 张 GPU 上启动 Llama 3 8B 模型的训练：

MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

多节点训练示例 (Slurm 环境)

若在集群环境（如 Slurm）运行，需修改提交脚本 multinode_trainer.slurm 中的节点数和 GPU 数，然后提交任务：

1. 编辑脚本设置节点数：

   #SBATCH --ntasks=2
   #SBATCH --nodes=2
   

2. 启动训练（确保 nnodes 与脚本设置一致）：

   srun torchrun --nnodes 2 ...
   

关键功能配置

所有训练选项均可通过 CLI 参数灵活配置：

• 混合精度/量化: 支持 Float8 和 MXFP8 (需 Blackwell GPU)。
• 并行策略: 自动组合 FSDP2、张量并行 (TP)、流水线并行 (PP) 和上下文并行 (CP)。
• 监控: 支持 Tensorboard 和 Weights & Biases 记录 Loss、显存、吞吐量等指标。

详细的高级配置和自定义扩展请参考项目目录下的 docs/ 文件夹及 torchtitan/models/llama3/config_registry.py。