torchtitan 正处于积极开发中。要使用 torchtitan 的最新功能，我们建议使用最新的 PyTorch nightly 版本。

最新动态

• [2025/11] AMD 发布了针对 AMD GPU 的 torchtitan 优化分支。
• [2025/10] 我们发布了 torchtitan v0.2.0。
• [2025/10] SkyPilot 现在支持 torchtitan！请参阅此处的教程。
• [2025/07] 我们发布了关于如何将模型添加到 torchtitan 的说明。
• [2025/04] 我们的论文被 ICLR 2025 接收。
• [2024/12] 关于 torchtitan 的 GPU MODE 讲座。
• [2024/07] 在 PyTorch Conference 2024 上的演示。

概述

torchtitan 是一个 PyTorch 原生平台，专为生成式 AI 模型的快速实验和大规模训练而设计。作为 PyTorch 原生扩展技术的最小化干净室实现，torchtitan 为开发者提供了灵活的基础。借助 torchtitan 扩展点，可以轻松创建针对特定需求的自定义扩展。

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

构建 torchtitan 的指导原则

• 旨在易于理解、使用和扩展，以适应不同的训练目的。
• 在应用多维并行时，对模型代码的改动最小化。
• 倾向于保持干净、最小的代码库，同时提供基本的可复用/可替换组件。

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

贡献

我们期待您的贡献！

• 为了加速对 torchtitan 的贡献及其相关创新，我们维护了一个 experiments 文件夹。新想法应从此处开始。要做出贡献，请遵循 experiments guidelines。
• 对于核心部分的修复和贡献，请遵循这些 guidelines。

Llama 3.1 训练

可用的主要功能

1. Multi-dimensional composable parallelisms
   • FSDP2 with per-parameter sharding
   • Tensor Parallel (including async TP)
   • Pipeline Parallel
   • Context Parallel
2. Meta device initialization
3. Per-op selective and full activation checkpointing
4. Distributed checkpointing (including async checkpointing)
   • Interoperable checkpoints which can be loaded directly into torchtune for fine-tuning
5. torch.compile support
6. Float8 support (how-to)
7. MXFP8 training for dense and MoE models on Blackwell GPUs.
8. Supervised Fine-Tuning (SFT) with chat-formatted datasets
9. DDP and HSDP
10. TorchFT integration
11. Checkpointable data-loading, with the C4 dataset pre-configured (144M entries) and support for custom datasets
12. Gradient accumulation, enabled by giving an additional --training.global_batch_size argument on the CLI
13. Flexible learning rate scheduler (warmup-stable-decay)
14. BF16 optimizer states for reduced memory usage
15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via Tensorboard or Weights & Biases
16. Debugging tools including CPU/GPU profiling, memory profiling, Flight Recorder, etc.
17. All options easily configured via Python config registry with --module and --config CLI flags
18. Structured logging: per-rank trace of key training phases; (see torchtitan/observability/structured_logger/README.md)
19. Helper scripts to
    • download tokenizers from Hugging Face

• 将原始 Llama 3 检查点转换为预期的 DCP 格式
  • 在不实例化模型的情况下估算 FSDP/HSDP 内存使用量
  • 使用 Tensor Parallel 运行分布式推理

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

深入代码

您可能想了解模型是如何定义的，或者并行技术是如何应用的。对于引导式浏览，请先查看以下文件：

• torchtitan/train.py - 主训练循环和高层设置代码
• torchtitan/models/llama3/model.py - Llama 3.1 模型定义
• torchtitan/models/llama3/parallelize.py - 用于将 Data Parallel、Tensor Parallel、activation checkpointing 和 torch.compile 应用于模型的辅助函数
• torchtitan/distributed/pipeline_parallel.py - 用于将 Pipeline Parallel 应用于模型的辅助函数
• torchtitan/components/checkpoint.py - 用于保存/加载分布式检查点的工具
• torchtitan/components/quantization/float8.py - 用于应用 Float8 技术的工具

安装

可以直接运行源代码，或者从 nightly build 或稳定版本安装 torchtitan。

从源代码

此方法需要 PyTorch 的 nightly build，或从 源代码 构建的最新 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 夜间构建版本时，需要 torchdata 的夜间构建版本。请按照上述所示从夜间构建索引中安装。

注意： 您可以直接从源代码树运行。如果需要从其他位置以包的形式导入 torchtitan，请在不重新解析依赖项的情况下以可编辑模式安装：pip install -e . --no-deps。

夜间构建

此方法需要 PyTorch 的夜间构建版本。您可以将 cu130 替换为其他版本的 cuda 或 AMD GPU（例如 rocm6.3）。

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

稳定版本

可以通过 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 分词器下载到本地机器。

# Get your HF token from https://huggingface.co/settings/tokens

# Llama 3.1 tokenizer
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。

Citation

我们提供了对 torchtitan 中可用的并行化和优化的详细分析，以及关于何时使用各种技术的总结建议。

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

源代码在 BSD 3 license 下提供，但您可能还有其他法律义务，这些义务可能约束您对本仓库中链接的其他内容的使用，例如第三方数据和模型的许可证或服务条款。