Chapter 01
项目背景与定位
1.1 一句话定位
nanotron 是 HuggingFace 团队 2023 年开源的、主打"小到能完整读完一遍"的 3D 并行预训练库。整个仓库的核心代码加起来才几千行,但已经包含张量并行 + 流水线并行 + 数据并行 + ZeRO-1 风格的 Distributed Optimizer,足以支撑 HF 内部 SmolLM、Idefics-2 等模型的训练。
它是 HuggingFace 在 Megatron-LM 之外的"自家选项"——更轻、更 PyTorch 原生、用 dataclass 驱动配置,特别适合做教学和魔改。
1.2 为什么会有 nanotron
HuggingFace 内部本来也用 Megatron-LM 和 DeepSpeed,但在大量训练实验中遇到两类痛点:
- Megatron-LM 的代码深度耦合 NVIDIA 自家 kernel,外人想加新算子或新并行策略要花很多力气;
- DeepSpeed 的 ZeRO 用得很顺,但 TP/PP 一直不是 first-class citizen。
HF 内部需要一份能完全 mastery 的代码:要能跑出 80%+ 的 MFU,要能在几个月内给一个新员工讲清楚每一行的目的。于是有了 nanotron。它的设计目标非常明确:"如果一份框架小到你能在一个周末读完,那它就值得长期维护。"
1.3 它在 HF 生态里的位置
| 名字 | HF 里的角色 |
|---|---|
| transformers | 模型仓库 / 推理 / 训练(HF Trainer) |
| accelerate | 让 PyTorch 代码"几行变分布式" |
| trl | 在 transformers 基础上做 SFT/DPO/PPO 训练 |
| peft | LoRA / QLoRA 等参数高效微调算法 |
| nanotron | 从零预训练大模型时使用 |
| datatrove | 数据清洗和 token 化流水线 |
简单说:当 HF 自己要从零训一个模型时,他们用 nanotron + datatrove;当用户拿别人的模型来微调时,用 transformers + trl + peft + accelerate。
1.4 设计哲学:少即是多
- 没有 monkey patch,没有 fused kernel 黑魔法。
nn.Linear就是nn.Linear,FlashAttention 通过参数选用; - 用 Python
@dataclass定义所有配置,运行时通过 YAML 文件填充; - ParallelContext 是唯一的"并行状态"对象,所有 group 信息从这里取;
- 代码中只有 3D 并行。没有 ZeRO-3,只有 ZeRO-1 风格的 Distributed Optimizer。HF 团队认为对他们的训练 workload,3D 并行 + ZeRO-1 已经够用。
代价就是:功能集合相对小。它不像 DeepSpeed 那样有 offload / NVMe / MoE 一大堆选项,也不像 Megatron 那样 FP8 / FlashAttention-3 全开。但对它服务的场景已经够了。
1.5 它适合用来做什么、不适合做什么
| 判断 | 典型场景 |
|---|---|
| ✅ 适合 | 学习"一个真实能跑起来的 3D 并行框架到底长什么样" |
| 学术研究 / 论文复现,需要 fork 一份能改的代码 | |
| HF 生态深度用户,要继续预训 SmolLM 风格的小到中型模型 | |
| 教学:给学生讲分布式训练时,这份代码比 Megatron 友好得多 | |
| ❌ 不适合 | 万亿模型 + 极致 MFU 的工业训练 |
| 想用 ZeRO-3 / offload / 异构内存等更激进显存优化 | |
| 完整 SFT/DPO 流水线(应该用 trl 或 LLaMA-Factory) |
1.6 本书将怎么带你走完这趟旅程
nanotron 因为代码量小,本书会尽量贴着代码讲:
- 第 2-3 章装环境 + 跑一个 tiny GPT;
- 第 4 章讲清楚
ParallelContext与 process group 的初始化; - 第 5 章逐个对照源码看 TP/PP/DP 的实现;
- 第 6-7 章讲 Distributed Optimizer 与配置系统;
- 第 8-10 章数据 / 模型 / 训练循环;
- 第 11-12 章调试和 checkpoint 转换;
- 第 13 章给一张"5000 行代码通读路线",是本书的特色——读完这一章你就能号称自己"读过一个完整的预训练框架"。