pytorch/torchtitan

torchtitan – 基于 PyTorch 的大规模生成式 AI 模型训练平台

是什么 – torchtitan 是由 PyTorch 团队开发的开源库，可让你使用原生 PyTorch 原语，以可扩展的方式训练 Llama 3.1 等大型语言模型（LLM）。它整合了多种分布式训练技术（FSDP2、张量并行、流水线并行、上下文并行等），并搭配检查点、量化和日志记录等实用工具，通过简洁、可扩展的 Python API 提供服务。

为何重要 – 训练现代 LLM（8 B 到 405 B 参数）通常需要复杂的自定义 CUDA 内核、第三方库和重型编排工具栈。torchtitan 旨在通过提供一个最小代码库，直接与最新的 PyTorch 夜间版本协同工作，降低这种复杂性。研究人员无需重写模型代码即可实验新的并行策略，而生产团队则可使用同一库从原型快速过渡到多节点训练。

核心功能（如 README 所列）

类别	功能
并行处理	• 多维可组合并行：FSDP2（按参数分片）、张量并行（含异步 TP）、流水线并行（零气泡）、超长序列用的上下文并行。
内存与速度优化	• 元设备模型初始化、选择性/完整激活检查点、分布式异步检查点、BF16 优化器状态、Float8 和 MXFP8 量化（Blackwell GPU）。
训练工作流	• 支持 `torch.compile`，通过令牌数标志实现梯度累积，灵活的预热-稳定-衰减型 LR 调度器，内置 C4 预配置数据加载器，支持自定义数据集插件。
可观测性	• TensorBoard / Weights & Biases 指标（损失、内存、吞吐量、MFU），每秩结构化日志，性能分析工具（CPU/GPU，飞行记录器）。
互操作性	• 分布式检查点格式（DCP）可直接加载至 `torchtune` 用于微调，支持 TorchFT 集成，提供下载分词器和检查点转换的辅助脚本。
可扩展性	• 文档完善的扩展点，`experiments` 文件夹支持快速原型（如 TitanRL RL 堆栈），添加新模型的清晰指南。
支持硬件	• NVIDIA GPU（已测试至 H100，支持 512 GPU 运行），AMD ROCm 分支，兼容任意 CUDA/ROCm 版本的夜间构建。

典型工作流

安装 – pip install torchtitan（或使用夜间 PyTorch + 夜间 torchtitan 轮子）。可在单节点或 Slurm/ParallelCluster 集群上开箱即用。
准备数据 – 使用内置 C4 加载器，或在配置中指向自定义数据集。
选择模型 – Llama 3.1（8 B、70 B、405 B）已内置；添加新模型请参考 torchtitan/models/README.md 中的简明指南。
配置 – 所有选项通过 Python 配置文件和 CLI 标志（--module, --config）暴露。可自由启用/禁用任何并行、检查点或量化选项。
启动 – 在 1～N 个节点上运行 torchrun（或提供的 run_train.sh）。torchtitan 会自动组合请求的并行策略并开始训练。
监控 – 指标在 TensorBoard/W&B 中显示；日志包含用于调试的每秩追踪信息。
检查点与微调 – 保存的 DCP 检查点可直接由 torchtune 加载，用于下游 SFT 或 RLHF。

快速入门（示例）

# 1️⃣ 安装夜间 PyTorch + torchtitan 轮子
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130

# 2️⃣ 下载 Llama‑3.1 分词器
python scripts/download_hf_assets.py \
    --repo_id meta-llama/Llama-3.1-8B \
    --assets tokenizer \
    --hf_token <your‑HF‑token>

# 3️⃣ 启动 8 GPU 运行（单节点）
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

CONFIG=llama3_8b 文件位于 torchtitan/config/，已启用 FSDP2 + 张量并行 + 激活检查点等配置。

适合谁使用？

研究人员：探索新扩展定律、并行组合或量化技术者。其干净室实现使底层代码易于阅读和修改。
工程师：构建生产级 LLM 预训练流水线，希望留在 PyTorch 生态系统内，避免外部专有堆栈者。
学生：学习分布式训练概念者。experiments 文件夹包含可运行的演示（如 TitanRL），展示同一模型代码如何用于 RL 训练和推理。

社区与支持

论坛 – 专用 PyTorch 论坛分类（pytorch‑forum → distributed → torchtitan）。
论文 – ICLR 2025 论文 《TorchTitan: 用于生产就绪 LLM 预训练的一站式 PyTorch 原生解决方案》（arXiv 2410.06511）。
贡献 – 核心变更和实验性想法的贡献指南见 CONTRIBUTING.md 和 torchtitan/experiments/README.md。
第三方分叉 – 存在 AMD 优化分叉（AMD‑AGI/torchtitan‑amd）。

许可证

torchtitan 采用 BSD‑3‑Clause 许可证发布。代码本身为宽松许可，但用户必须遵守所下载外部数据或模型权重的许可证（例如 Meta‑Llama 分词器和检查点）。

总结 – torchtitan 提供了一个生产就绪、仅依赖 PyTorch 的大规模 LLM 训练栈，强调代码可读性、可扩展性和最新的分布式训练研究。若需在不引入重型外部生态的情况下训练或实验大规模生成模型，此库是最佳起点。