pytorch/torchtitan

`torchtitan` 正在密集开发中。为了使用 `torchtitan` 的最新功能，我们建议使用最新的 PyTorch nightly 版本。 ## 最新动态 - [2025/11] AMD 发布了针对 AMD GPU 优化的 `torchtitan` [分支](https://github.com/AMD-AGI/torchtitan-amd/tree/main)。 - [2025/10] 我们发布了 `torchtitan` [v0.2.0](https://github.com/pytorch/torchtitan/releases)。 - [2025/10] SkyPilot 现已支持 `torchtitan`！请查看[此处的](https://docs.skypilot.co/en/latest/examples/training/torchtitan.html)教程。 - [2025/07] 我们发布了关于如何向 `torchtitan` 添加模型的[说明](/torchtitan/models/README.md)。 - [2025/04] 我们的论文被 [ICLR 2025](https://iclr.cc/virtual/2025/poster/29620) 收录。 - [2024/12] 关于 torchtitan 的 GPU MODE [讲座](https://www.youtube.com/watch?v=VYWRjcUqW6w)。 - [2024/07] 在 PyTorch Conference 2024 上的[演讲](https://pytorch2024.sched.com/event/1fHn3)。 ## 概述 `torchtitan` 是一个 PyTorch 原生平台，专为生成式 AI 模型的**快速实验和大规模训练**而设计。作为 PyTorch 原生扩展技术的极简纯净实现，`torchtitan` 为开发者提供了一个灵活的基础，便于在此基础上进行构建。借助 `torchtitan` 的[扩展点](docs/extension.md)，您可以轻松创建满足特定需求的自定义扩展。 我们的使命是通过赋能研究人员和开发者探索新的模型架构和基础设施技术，加速生成式 AI 领域的创新。 构建 `torchtitan` 时的指导原则 * 设计为易于理解、使用，并能针对不同的训练目的进行扩展。 * 在应用多维并行时，尽量减少对模型代码的修改。 * 倾向于简洁、极简的代码库，同时提供基本的可重用/可替换组件。 `torchtitan` 一直通过支持预训练各种规模的 Llama 3.1 LLM，来展示 PyTorch 最新的分布式训练特性。 ## Llama 3.1 训练 ### 可用的主要特性 1. 多维可组合并行 - 带有按参数分片的 [FSDP2](docs/fsdp.md) - [Tensor Parallel](https://pytorch.org/docs/stable/distributed.tensor.parallel.html)（包含[异步 TP](https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487)） - [Pipeline Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420) - [Context Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) 2. [Meta device](https://pytorch.org/docs/stable/meta.html) 初始化 3. 针对算子的选择性和完全 activation checkpointing 4. [分布式 checkpointing](https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250)（包含异步 checkpointing） - [可互操作的 checkpoint](docs/checkpoint.md)，可以直接加载到 [`torchtune`](https://github.com/pytorch/torchtune) 中进行 fine-tuning 5. 支持 `torch.compile` 6. 支持 [Float8](https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323)（[指南](torchtitan/components/quantization/float8.md)） 7. 在 Blackwell GPU 上支持[针对稠密和 MoE 模型的 MXFP8 训练](torchtitan/components/quantization/mxfp8.md)。 8. 使用聊天格式数据集进行监督微调 (SFT) 9. DDP 和 HSDP 10. 集成 [TorchFT](https://github.com/pytorch/torchft) 11. 可进行 checkpoint 的数据加载，预配置了 C4 数据集（1.44 亿条目），并支持[自定义数据集](docs/datasets.md) 12. 梯度累积，通过在 CLI 中提供额外的 `--training.global_batch_size` 参数启用 13. 灵活的学习率调度器 (warmup-stable-decay) 14. [BF16 optimizer states](docs/bf16_optimizer_states.md)，用于减少内存使用 15. Loss、GPU 内存、吞吐量 (tokens/sec)、TFLOPs 和 MFU 通过 [Tensorboard 或 Weights & Biases](/docs/metrics.md) 显示和记录 16. [调试工具](docs/debugging.md)，包括 CPU/GPU 性能分析、内存分析、Flight Recorder 等 17. 所有选项都可以通过 [Python 配置注册表](torchtitan/models/llama3/config_registry.py) 以及 `--module` 和 `--config` CLI 标志轻松配置 18. 结构化日志：记录关键训练阶段的每个 rank 的 trace；（参见 [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md)） 19. [辅助脚本](scripts/) 用于 - 从 Hugging Face 下载 tokenizer - 将原始的 Llama 3 checkpoint 转换为预期的 DCP 格式 - 在不实例化模型的情况下估算 FSDP/HSDP 内存使用量 - 使用 Tensor Parallel 运行分布式推理 我们报告了在多达 512 个 GPU 上的[性能](benchmarks/llama3_h100_202412_torchtitan.md)，并验证了各种技术的 [loss 收敛](docs/converging.md) 正确性。 ### 深入代码 您可能想了解模型是如何定义的，或者并行技术是如何应用的。如需引导式浏览，请先查看以下文件： * [torchtitan/train.py](torchtitan/train.py) - 主训练循环和高层设置代码 * [torchtitan/models/llama3/model.py](torchtitan/models/llama3/model.py) - Llama 3.1 模型定义 * [torchtitan/models/llama3/parallelize.py](torchtitan/models/llama3/parallelize.py) - 用于将 Data Parallel、Tensor Parallel、activation checkpointing 和 `torch.compile` 应用于模型的辅助工具 * [torchtitan/distributed/pipeline_parallel.py](torchtitan/distributed/pipeline_parallel.py) - 用于将 Pipeline Parallel 应用于模型的辅助工具 * [torchtitan/components/checkpoint.py](torchtitan/components/checkpoint.py) - 用于保存/加载分布式 checkpoint 的工具 * [torchtitan/components/quantization/float8.py](torchtitan/components/quantization/float8.py) - 用于应用 Float8 技术的工具 ## 安装 您可以直接运行源代码，或者从 nightly build 或稳定版本安装 `torchtitan`。 ### 从源码安装 此方法需要 PyTorch 的 nightly build，或者[从源码](https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source)构建的最新 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 ``` ### Nightly build 此方法需要 PyTorch 的 nightly build。您可以将 `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` [稳定版本](https://github.com/pytorch/torchtitan/releases)。 ``` pip install torchtitan ``` ``` conda install conda-forge::torchtitan ``` 请注意，每个稳定版本都会固定 `torch` 和 `torchao` 的 nightly 版本。请参阅 [release.md](docs/release.md) 了解更多详情。 ### 下载 tokenizer `torchtitan` 目前支持开箱即用地训练 Llama 3.1 (8B, 70B, 405B)。要开始训练这些模型，我们需要下载 tokenizer。请按照官方 [meta-llama](https://huggingface.co/meta-llama/Llama-3.1-8B) 仓库上的说明进行操作，确保您拥有对 Llama 模型权重的访问权限。 确认访问权限后，您可以运行以下命令将 Llama 3.1 tokenizer 下载到本地计算机。 ``` # 从 https://huggingface.co/settings/tokens 获取你的 HF token # 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`。 ## 引用 我们详细探讨了 `torchtitan` 中可用的并行技术和优化，并总结了何时使用各种技术的建议。 [TorchTitan：用于生产级 LLM 预训练的一站式 PyTorch 原生解决方案](https://openreview.net/forum?id=SFN6Wm7YBI) ``` @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 许可证](./LICENSE)下提供，但是，您可能还需要遵守其他法律义务，这些义务约束了您对本仓库中链接的其他内容的使用，例如第三方数据和模型的许可证或服务条款。

标签：PyTorch, Yelp, 凭据扫描, 分布式训练, 模型训练, 深度学习, 生成式AI, 索引, 逆向工具