pytorch/torchtitan
A PyTorch native platform for training generative AI models
torchtitan – 面向大规模生成式AI训练的PyTorch原生平台
它是什么
torchtitan 是PyTorch团队的开源库,让你能够完全在PyTorch中训练、微调,并对超大规模语言模型(LLM)和多模态模型进行强化学习(RL)。代码库刻意保持小巧且可读,并将整个训练流水线表达为Python“配方”——一个返回配置对象的函数。通过切换配方,你可以改变模型大小、数据、并行策略,甚至无需重写模型即可从预训练切换到RL。
核心思想及其重要性
| 思想 | 功能 | 重要性 |
|---|---|---|
| 可组合并行 | 模型声明应如何分片;库通过DeviceMesh和DTensor堆叠全分片数据并行(FSDP2)、张量并行(TP)、流水线并行(PP)、上下文并行(CP)和专家并行(EP)。 |
让研究人员在同一代码上尝试任何并行组合,无需自定义内核即可从单个GPU扩展到数百个GPU。 |
| 从预训练到RL的一栈式 | 相同的模型定义和内核用于经典语言模型预训练、监督微调(SFT),以及驱动vLLM生成的TitanRL强化学习循环。 | 保证训练器与推理之间逐位可复现,简化基于RL的对齐或策略梯度工作的调试。 |
| 前沿模型支持 | 提供Llama 3、Qwen 3、DeepSeek V3、FLUX、Kimi、Muse Glimmer等的现成配方,包括低精度格式(MXFP8、NVFP4)和MoE(Dist-MoE)。 | 研究者无需重新实现架构或量化技巧即可在最先进的模型上开始实验。 |
| Python优先配置 | 配方只是一个返回配置对象的Python函数;你编辑函数而不是传递一长串CLI标志。 | 让实验可复现且易于版本控制;你可以与同事共享配方文件,他们获得完全相同的运行结果。 |
| 广泛的可观测性 | 内置指标(损失、TFLOPs、MFU)、TensorBoard/W&B日志记录、性能分析工具、确定性模式、损失比较脚本以及SDC重放。 | 帮助及早发现性能瓶颈或数值错误,这在数百个GPU规模下至关重要。 |
今天你可以用它做什么
- 训练新LLM – 选择
llama3_8b等配方,在8个GPU上启动./run_train.sh;库会自动处理分片、检查点和数据加载。 - 使用LoRA或SFT微调 – 相同的配方系统允许你添加LoRA适配器或聊天数据集微调步骤。
- 运行基于RL的对齐 – 使用
TitanRL栈运行PPO风格或更新的GRPO/DAPO损失,重用训练时的确切模型权重。 - 实验低精度训练 – 通过配置标志启用MXFP8或NVFP4;库附带必要的
torchao内核。 - 扩展到数百个GPU – 文档化的基准测试显示在H100上最多512个GPU的运行,具有自动分布式检查点和异步保存。
谁应该使用它
- 研究实验室 – 需要干净、可入侵的代码库来探索新的并行或RL思想。
- 工程师 – 构建生产级LLM训练流水线,希望留在PyTorch生态系统中(无需单独的C++/CUDA框架)。
- 学生 – 学习大规模训练 – 配方是展示如何组合FSDP、TP和MoE的可读示例。
快速入门(摘要)
- 安装 – 克隆仓库,安装PyTorch nightly和
torchtitan包(或使用预构建的pip/conda轮子)。 - 运行调试作业 –
NGPU=1 ./run_train.sh训练一个微小的Llama-3调试模型10步。 - 训练真实模型 – 从Hugging Face下载分词器,设置
MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8b,并在8个GPU上启动。 - 自定义 – 编写自己的配方函数,调整并行度、数据混合或添加LoRA。
文档与社区
- 文档分为多个主题部分(配置、模型、分布式、数据、检查点、RL、调试、基准测试)。每个文件夹包含带有使用示例的
README.md。 - CI徽章显示在CPU、NVIDIA(包括H100/B200)、AMD ROCm和RL工作负载上的积极测试。
- 贡献 – 指南鼓励在
experiments/文件夹中进行新实验,并通过拉取请求进行核心改进。 - 论文与引用 – ICLR-2025论文详述了设计选择;如果你发表基于torchtitan的作品,请引用。
以上信息均直接取自仓库的README;未添加任何外部假设。
相关
- 项目
- 项目
- 项目
- 项目
- 项目