Chapter 01
项目背景与定位
1.1 一句话定位
TorchTitan 是 PyTorch 团队（Meta + 社区）2024 年开源的、作为 PyTorch 官方"如何用原生能力训大模型"参考实现的代码库。它的目标不是替代 Megatron-LM 或 DeepSpeed，而是用一份干净的 Llama 训练脚本告诉所有人：PyTorch 自带的 FSDP2 + DTensor + TP + PP 已经足够训出 SOTA 大模型，不需要任何第三方 patch。
如果你之前一直疑惑"FSDP 到底能不能扛得起千亿模型"、"DTensor 到底有什么用"——TorchTitan 就是 PyTorch 给出的正式答案。
1.2 为什么会有 TorchTitan
PyTorch 团队为什么要"亲自下场"再写一个训练框架？三个动机：
- DeepSpeed / Megatron 的存在感太强，让外界感觉 PyTorch 原生能力不行。事实上 FSDP（PyTorch 1.11+）和 DTensor（2.0+）都已经具备相应能力，但需要一份"高质量公开示范"。
- PyTorch 内部新特性需要 dogfooding：FSDP2、torch.compile、float8 training、async TP……
- 统一参考：每次 issue 里问"应该怎么用 PyTorch 训 Llama"，需要一个标准答案能 link 过去。
于是有了 TorchTitan：仓库小（万行级别）、依赖极少（只要 PyTorch + tokenizer）、focus 单一（Llama 系列预训练为主），把 PyTorch 团队认为"2024 年应该这样写训练代码"的所有 best practice 装进去。
1.3 它和 Megatron / DeepSpeed 的关系
| 维度 | TorchTitan | Megatron-LM | DeepSpeed |
|---|---|---|---|
| 维护方 | PyTorch 团队 | NVIDIA | Microsoft |
| 基础抽象 | DTensor / DeviceMesh / FSDP2 | NCCL + 自家 parallel_state | ZeRO + 自家通信 |
| 与 PyTorch 关系 | "我就是 PyTorch" | 在 PyTorch 之上 + 大量自家 kernel | 在 PyTorch 之上 + 自家 op |
| 代码量 | 小（万行级） | 中（10w 行级） | 大（10w 行级） |
| 学习曲线 | 中（依赖你熟 PyTorch 新特性） | 陡 | 中 |
| 推荐用法 | 改 toml + 写自己的 model | 改 args + 复用脚本 | JSON 配 ZeRO + Trainer |
经验法则：
- 想"长期跟着 PyTorch 走、不引第三方 patch、有能力 follow PyTorch 新特性" → TorchTitan；
- 想"立刻就能跑出 SOTA 性能、有 NVIDIA 集群" → Megatron-LM；
- 想"已有 HF Trainer 代码、想加分布式" → DeepSpeed。
1.4 设计哲学：原生 PyTorch + 极简代码
- 拒绝任何 monkey patch。它只用
torch.distributed.fsdp.fully_shard、torch.distributed.tensor.parallel.parallelize_module这些公开 API； - 配置极简。一份
train_configs/llama3_8b.toml几十行就能描述整个训练； - 模型代码"看上去就是普通 PyTorch"；
- 追随 PyTorch 路线图。每次 PyTorch 大版本发布的新分布式能力，TorchTitan 都会第一时间 dogfood。
代价是：你必须接受 nightly PyTorch / 较新的 CUDA，否则很多功能可能要么没法用、要么有 bug。
1.5 它适合用来做什么、不适合做什么
| 判断 | 典型场景 |
|---|---|
| ✅ 适合 | 想要训 Llama 系列、用 PyTorch 原生能力的项目 |
| 学习"PyTorch 官方推荐的分布式写法" | |
| 有研究新算子 / 自定义架构的需求，希望底座尽量薄 | |
| 想做 FP8 训练、想跟 PyTorch 团队同步前沿能力 | |
| ❌ 不适合 | 生产环境要求"用稳定版 PyTorch"，TorchTitan 经常依赖 nightly |
| 完整的 SFT/DPO/PPO 工业流水线 —— 它是预训练 demo | |
| 想要丰富 plugin、不想自己写模型代码 —— ColossalAI / LLaMA-Factory 更合适 |
1.6 本书将怎么带你走完这趟旅程
- 第 2-3 章带你装好 PyTorch nightly、跑通 8 卡 Llama-3 8B 训练；
- 第 4 章讲透
DTensor+DeviceMesh，这是后面所有章节的基础抽象； - 第 5-6 章深入 FSDP2 与 PyTorch 原生 TP/PP；
- 第 7-8 章拆解 toml 配置 + 整份
train.py； - 第 9 章讲 PyTorch 新特性怎么打开（compile / async TP / float8）；
- 第 10-12 章实战 + 调优 + 自定义架构；
- 第 13 章给一张源码导读图，并指出 PyTorch 分布式后续会往哪走。