前言:从 Lua 边缘到 AI 心脏
“PyTorch 的设计核心是:让普通的 Python 程序员能够写出科研级别的模型,而不需要先学一门新语言。”
—— Soumith Chintala,PyTorch 项目创始人之一
一切的起点:纽约大学一台老式工作站
故事要从 2002 年纽约大学 Yann LeCun 的实验室讲起。
那一年,他们正在用一种叫 Lush 的 Lisp 方言做神经网络研究。Lush 是 LeCun 自己参与设计的语言,把神经网络当作一等公民:张量、卷积、反向传播都内嵌在语法层。但它也带着 Lisp 与生俱来的孤独 —— 整个工业界都在用 C++ 与 Java,没有公司愿意为一种小众语言养工程师,研究生们在毕业之后再也找不到能继续用 Lush 写代码的地方。
LeCun 的学生 Ronan Collobert 决定换一个底座。2007 年前后,他和 Koray Kavukcuoglu、Clement Farabet 把 Lush 里关于张量和神经网络的核心思想搬到了一种当时小有名气的脚本语言 Lua 上 —— 这就是 Torch7。
Torch7 用 Lua 写前端,C 写底层,配合一个叫 TH(TorcH 的缩写)的张量库做底层 BLAS 和 BLAS-like 的运算。它在学术界站稳了脚跟:DeepMind 在用、Facebook AI Research 在用、Twitter Cortex 在用。但 Lua 和 Python 之间的鸿沟,注定 Torch7 永远只能是少数派。每招进一个新研究员,团队都要先教会他写 Lua —— 这件事在 Python 一统机器学习的 2014 年,已经变得不可持续。
2016 年,Soumith Chintala、Adam Paszke、Sam Gross 等人在 Facebook AI Research 启动了一个新项目,把 Torch7 的张量库和自动微分引擎整体搬到 Python 上。一年之后,2017 年 1 月 18 日,PyTorch 0.1.6 alpha 0 在 GitHub 上公开发布。
这就是今天每一篇顶会论文、每一份大模型权重、每一个国内外 AI 实验室默认背后的那个项目的开端。
Caffe2 合并:把推理框架”喂”进 PyTorch
PyTorch 0.x 时代有一个明显的短板:它只擅长研究,不擅长部署。Facebook 内部当时还有另一个深度学习框架叫 Caffe2,由贾扬清(贾扬清后来从 Facebook 去了阿里、又去了 LeptonAI)从 Caffe 演化而来,专门为生产推理优化 —— C++ 内核、轻量级、能跑在手机上。
2018 年 5 月的 F8 开发者大会上,Facebook 宣布了一个重磅决定:把 Caffe2 整体合并进 PyTorch。Caffe2 的张量库、运行时与 ONNX 模型互转能力,全部并入 PyTorch 仓库。这个决定带来了一个深远的影响 —— PyTorch 从此既能做研究也能做部署,学术和工业不必再”双框架”切换。
这次合并在源码上留下的痕迹今天还能看到:
caffe2/目录长期保留在 PyTorch 主仓库里(直到 v2.5 才被彻底清除)- 早期的 ATen 张量后端继承了 Caffe2 的
Tensor接口 (caffe2::Tensor/c10::TensorImpl) c10这个名字本身就是 Caffe2 与 A Ten(“a Tensor”)的混搭缩写
PyTorch 1.0(2018 年 12 月)就是这次合并后的第一个稳定版本。理解 PyTorch 的”双底色” —— 研究友好 + 工业可部署 —— 就要回到 Caffe2 合并这个时间点。
PyTorch 大版本演进路线图
为了让你对 PyTorch 各代版本的”性格”有个直观认识,把关键里程碑列出来:
| 版本 | 时间 | 核心变化 | 代表性能力 |
|---|---|---|---|
| 0.1 | 2017-01 | 首次公开 | 动态图 + autograd |
| 0.4 | 2018-04 | 合并 Tensor 与 Variable | 用户不再区分两者 |
| 1.0 | 2018-12 | Caffe2 合并 + JIT/TorchScript | 同框架研究 → 部署 |
| 1.3 | 2019-10 | 量化、移动端、命名张量 | 工业部署能力补齐 |
| 1.5 | 2020-04 | C++ 前端稳定、分布式 RPC | 多机训练原生支持 |
| 1.8 | 2021-03 | TorchElastic、Pipeline Parallelism | 大模型训练抽象层 |
| 1.10 | 2021-10 | DDP no_sync、CUDA Graph 实验性 | 性能基建持续打磨 |
| 1.12 | 2022-06 | functorch、FSDP 引入 | 函数式 transform、ZeRO 风格 |
| 2.0 | 2023-03 | torch.compile 公开 | Dynamo + AOTAutograd + Inductor |
| 2.2 | 2024-01 | FlashAttention-2、scaled_dot_product_attention | LLM 训练加速 |
| 2.4 | 2024-07 | FSDP-2 (_fully_shard) | 新一代参数分片 |
| 2.6 | 2025-01 | PT2E 量化、Distributed Checkpoint 稳定 | 大规模训练工程化 |
| 2.8 | 2025-07 | Compiled Autograd 稳定 | 反向图也能被 Inductor 编译 |
| 2.11 | 2026-03 | 本书使用版本 | torch.compile 默认、FSDP-2 默认 |
最值得停下来看一眼的是 2.0 这个分水岭:在 2.0 之前,PyTorch 的故事是”如何把动态图做得更好用”;在 2.0 之后,故事变成了”如何在保留动态图体验的前提下,把性能拿到静态图水平”。这两条故事线在源码里留下的痕迹完全不同 —— 前者堆叠在 aten/ 与 torch/csrc/autograd/ 下,后者集中在 torch/_dynamo/、torch/_functorch/、torch/_inductor/ 三个目录。理解 PyTorch 源码结构,就要理解这条 2.0 分水岭。
与国内框架的对比
很多读者关心的另一个问题:国内的 PaddlePaddle、Mindspore、OneFlow 与 PyTorch 是什么关系?
简短的回答:
| 框架 | 主导方 | 与 PyTorch 的关系 |
|---|---|---|
| PaddlePaddle | 百度 | 早期类似 TF 1.x 静态图,后来推动动态图 + 静态图融合,API 上对 PyTorch 有大量兼容借鉴 |
| Mindspore | 华为 | 自研图模式 + Pynative 动态模式,长期主打”昇腾原生”。API 与 PyTorch 有差异 |
| OneFlow | 一流科技(已被字节并购) | 基于”全局视角并行” SBP(Split / Broadcast / Partial),API 接近 PyTorch,编译技术领先 |
国内框架的共同特点是 “PyTorch 兼容层 + 自研编译器后端 + 自研芯片支持” —— 用户写 PyTorch 风格的代码,框架在底层把它翻译到自家的运行时和加速卡上。理解 PyTorch 源码,相当于理解了国内框架们都在试图兼容的”事实标准”。这也是为什么本书的内容对所有 AI 工程师都有迁移价值,而不仅仅是 PyTorch 用户。
graph LR
Lush["Lush (NYU Lisp 方言)<br/>1990s-2000s"]
Torch7["Torch7 (Lua + C)<br/>2002-2017"]
PyTorch["PyTorch (Python + C++)<br/>2017-至今"]
Caffe2["Caffe2 (C++)<br/>Facebook 推理框架<br/>2017-2018"]
Lush --> Torch7
Torch7 --> PyTorch
Caffe2 -. "2018 合并" .-> PyTorch
style PyTorch fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
一个数字说明 PyTorch 的地位
如果今天你在 Papers With Code 上随便点开一份 NeurIPS / ICLR / ICML 的 best paper:
2024 年顶会论文使用的训练框架统计 (Papers With Code 公开数据):
PyTorch: ~92%
JAX: ~5%
TensorFlow: ~2%
其他 (Mindspore 等): ~1%
如果你打开 HuggingFace Hub,超过 80% 的开源模型权重以 PyTorch 格式(.bin / .safetensors)作为首发格式。Llama、Qwen、DeepSeek、Mistral、Gemma —— 每一个你在新闻里看到的开源大模型,背后都是一行 torch.save(model.state_dict(), 'pytorch_model.bin')。
如果你算一下硬件:
2025 年全球 GPU 训练算力规模:
NVIDIA H100/H200: ~500 万张
GB200 / B100: ~30 万张
其他 GPU 训练卡: ~200 万张
其中 PyTorch 占用份额: ~85%
对应每年的电费 + 折旧: 千亿美元量级
PyTorch 是当今 AI 算力上跑的字面意义上的第一大操作系统。每一秒钟,都有数百万张 GPU 在执行某段被 PyTorch 编排出来的计算图。每提升 1% 的训练效率,都是在节省国家级别的能耗。
理解 PyTorch 的内部机理,不是一项可有可无的修养;它是这个时代 AI 工程师的基本功。
一个被低估的事实:PyTorch 是”事实上”的训练协议
很多人以为 PyTorch 只是”一个框架”。但更准确的描述是:PyTorch 已经成为深度学习的事实上的协议层,类似于 HTTP 之于互联网、JVM 之于 Java 生态。
证据是:
- HuggingFace Transformers 库内部以 PyTorch 为第一公民。即便它支持 TensorFlow 与 JAX,新模型几乎都先用 PyTorch 写
- DeepSpeed(微软)、Megatron-LM(NVIDIA)、ColossalAI(潞晨科技)、InternEvo(上海 AI 实验室)—— 几乎所有大模型训练框架都建立在 PyTorch 之上
- vLLM、SGLang、TensorRT-LLM 这些主流推理引擎,输入端默认接受的是 PyTorch 格式的权重
- OpenAI Triton 编译器虽然独立,但绝大多数被用作
@torch.compile的后端 - 国产 AI 芯片(华为昇腾、寒武纪、壁仞、海光)几乎都在花大量精力做 PyTorch 兼容层
这意味着,理解 PyTorch 不再只是”会用一个工具”,而是 理解整个深度学习行业是怎么运转的。
为什么又一本 PyTorch 书
如果你搜 “PyTorch tutorial” 或 “PyTorch 教程”,能找到几千篇博客、几百小时视频、一份越来越厚的官方文档。它们大多数的套路是:
- 怎么用
nn.Linear搭一个 MLP - 怎么写
forward方法 - 怎么
loss.backward()然后optimizer.step() - 怎么
torch.save模型再加载推理 - 怎么把代码搬到 GPU
这些都有用,但合起来仍然只是 “使用手册” —— 告诉你按哪些按钮,却不告诉你按钮背后有什么线路。
工程师真正想知道的
而工程师真正想知道的是另一组问题:
- 一个
tensor.add(other)到底经过了多少层调度,最终落到哪个 CUDA kernel? .cuda()把张量搬到 GPU,到底是cudaMalloc还是从某个池子里捞?为什么连续创建一万个张量不会让显存爆掉?loss.backward()是怎么知道要从哪里开始反向传播的?反向图什么时候构建?- 多 GPU 的 DDP 怎么做到”几乎线性扩展”?环形 AllReduce 的环到底怎么转?
torch.compile是黑魔法吗?它究竟拦截了什么、生成了什么?为什么有时候它让代码快 2 倍、有时候反而慢?- FSDP 的 ZeRO-3 是怎么做到”只用 1/N 的显存训 N 倍大的模型”的?参数分片在前向时什么时候 unshard、反向时什么时候 reshard?
- 自定义算子要写哪些声明文件?为什么不能直接拿 C++ 写一个函数注册上去就完事?
- DataLoader 的多进程 worker 之间是怎么通信的?为什么有时候
num_workers=8反而比num_workers=4慢? - Profiler 给出的 “CUDA Total Time” 和 “CPU Total Time” 到底对应代码的哪一段?
- 当训练突然变慢,第一个该看的瓶颈是什么?第二个、第三个又在哪里?
这些”为什么”藏在源码注释里、PR 讨论里、issue 追踪里、PyTorch 开发者大会某场闪电演讲的最后一张幻灯片上。你要把它们拼起来,才能真正理解 PyTorch。
本书的定位
本书做的就是这件事 —— 把源码、设计文档、PR 讨论、benchmark 数据和真实工程踩坑串成可读的章节 —— 不是源码导读,是源码之后的思考。
graph LR
A["官方教程<br/>"怎么用""] --> B["社区博客<br/>"为什么这样用""]
B --> C["源码<br/>"到底是怎么做的""]
C --> D["本书<br/>"为什么这样设计""]
style A fill:#fef3c7,stroke:#f59e0b
style B fill:#fef3c7,stroke:#f59e0b
style C fill:#dbeafe,stroke:#3b82f6
style D fill:#dcfce7,stroke:#22c55e,stroke-width:2px
本书基于的源码版本
本书基于 PyTorch v2.11.0(2026-03 发布),对应 commit hash:
git tag: v2.11.0
commit: 70d99e998b4955e0049d13a98d77ae1b14db1f45
release: 2026-03-20
选择 v2.11 而非更早的 v2.0 / v2.4,是因为它是几条历史脉络汇合后的稳定形态:
- torch.compile 全面成熟:从 2.0 (2023) 引入到 2.11,Dynamo 与 Inductor 经历了 3 年的工业级打磨,之前的实验性 flag 大部分已经成为默认行为
- FSDP-2 默认推荐:
torch.distributed._composable.fsdp接口(FSDP-2)已经在 v2.4 之后成为新代码的默认推荐路径,旧的FullyShardedDataParallel类被标记为兼容层 - PT2E 量化稳定:
torch.ao.quantization.pt2e(基于 export + 编译器的量化)取代了旧的 FX graph mode 量化 - Distributed Checkpoint 标准化:
torch.distributed.checkpoint已经成为千卡级训练的事实保存格式 - C10D / NCCL 升级:进程组接口、ProcessGroupNCCL 的健康检查、DTensor 都进入稳定期
理解 v2.11,就是理解 PyTorch 在”训练框架应该长什么样”这个问题上的当前答案。
源码对照指南
我建议你在阅读时同步打开源码对照:
# 拉取本书使用的精确版本
git clone --depth=1 --branch v2.11.0 \
https://github.com/pytorch/pytorch.git pytorch-2.11
cd pytorch-2.11
# 五大核心路径
ls c10/ # 核心抽象:Device、DType、Allocator、intrusive_ptr
ls aten/src/ATen/ # 算子库:native_functions.yaml + 实现
ls torch/csrc/ # Python 与 C++ 的桥梁,autograd C++ 引擎在这里
ls torch/ # Python 前端:nn、optim、data、distributed
ls torch/_dynamo/ # TorchDynamo:CPython 帧拦截、字节码分析
ls torch/_inductor/ # TorchInductor:编译器后端,Triton 代码生成
ls torch/distributed/ # 分布式训练:DDP / FSDP / 进程组 / NCCL 集成
# 看代码生成出来的真实算子定义(很多人不知道这一步)
cat aten/src/ATen/native/native_functions.yaml | head -100
# 看 Dispatcher 注册了哪些 key
grep -r "TORCH_LIBRARY_IMPL" aten/src/ATen/native/ | head
引用标注规范
书中所有源码引用都标注了具体路径,形如:
aten/src/ATen/core/dispatch/Dispatcher.h:142 ← 精确行号引用
torch/_dynamo/symbolic_convert.py ← 模块级引用
c10/core/TensorImpl.h:380-410 ← 区间引用
遇到一个精确的行号引用,说明那一行代码足够关键,值得你停下来对照读一遍。
如果将来 PyTorch 升级到 v2.12 / v2.13,行号会漂移,但模块路径和符号名稳定性极高 —— 找不到具体行号时,按符号名搜一下基本能命中。
为什么是 PyTorch 而不是其他
在 PyTorch 之前,深度学习框架的赛场上至少有过六七个候选人:
| 框架 | 时代 | 强项 | 退场或衰落原因 |
|---|---|---|---|
| Theano | 2008-2017 | 自动微分先行者,蒙特利尔大学出品 | 编译慢、API 古怪、被 TensorFlow 完全替代 |
| Torch7 | 2002-2017 | 学术圈实力派 | Lua 语言生态太小 |
| Caffe | 2013-2017 | CV 模型部署利器 | 配置文件式声明,灵活性差 |
| Caffe2 | 2017-2018 | Facebook 自家推理框架 | 2018 与 PyTorch 合并 |
| MXNet | 2015-2023 | 亚马逊主推,多语言绑定 | 2023 进入 Apache Attic(即”项目坟场”) |
| TensorFlow | 2015-至今 | Google 出品,工业部署强 | 1.x 过于声明式,2.x 体验混乱,研究圈大量逃离 |
| JAX | 2018-至今 | 函数式 + XLA 编译 | API 抽象层级高,研究圈双语派系,工业落地慢 |
| PyTorch | 2017-至今 | 动态图 + Python 原生 + 庞大社区 | —— 当下事实上的赢家 |
动态图革命:PyTorch 真正赢在哪里
PyTorch 不是第一个支持自动微分的库,也不是第一个 Python 优先的库。它真正的胜负手是 动态图(define-by-run)。
什么叫动态图?看这段代码:
import torch
def model(x, depth):
h = x
for i in range(depth): # depth 在每一次调用时可以不同
if h.sum() > 0: # 控制流可以依赖中间结果
h = torch.relu(h @ W[i])
else:
h = torch.tanh(h @ W[i])
return h
# depth=3 跑出来一个图
y1 = model(x, 3)
# depth=5 又是另一个图
y2 = model(x, 5)
在 PyTorch 里,计算图是在 Python 解释器执行的过程中”边走边长”的。每次调用都构建一个新图,反向传播完就丢掉。Python 的 if / for / 异常都直接是图的控制流。
而 TensorFlow 1.x 走的是静态图(define-and-run):
# TF 1.x 的写法(已经被淘汰)
x = tf.placeholder(tf.float32)
W = tf.Variable(...)
y = tf.matmul(x, W) # 这一行不执行计算,只是在描述图
with tf.Session() as sess:
out = sess.run(y, {x: real_data}) # 真正的执行在这里
你必须先用 tf.placeholder / tf.Variable / tf.matmul 这些 API 把整张图描述出来,再用 Session.run() 喂数据真跑。控制流必须用 tf.cond / tf.while_loop 这些图原语,不能直接用 Python 的 if / for。
对工业部署,静态图有它的好处:可序列化、可优化、可静态分析。但对研究,静态图是噩梦:调试要 tf.print,动态形状要小心,新论文里那些”看一眼输入决定走哪条分支”的 trick 全部要重新建模成图原语。
PyTorch 用动态图把研究门槛打到了几乎为零:你写的就是普通 Python,pdb 能打断点,print 能立刻看到张量值。这一招让 PyTorch 在 2018-2020 年间几乎独占了学术界 —— 新论文用 PyTorch 实现,工业界为了跑这些论文不得不切到 PyTorch,进而 PyTorch 又积累了更多生态。这是一个完整的飞轮。
到 2019 年,TensorFlow 团队意识到不对,推出了 TF 2.0 + Eager Mode 想模仿动态图。但势已成 —— 研究者已经用 PyTorch 习惯了,工业界已经把推理流水线切换到了 PyTorch。PyTorch 赢在 timing:在动态图最被需要的时刻,把它做对了。
那 torch.compile 算什么?
讽刺的是,PyTorch 在赢得动态图之战后,又在 2022 年的 PyTorch 2.0 大会上推出了 torch.compile —— 本质上把动态图重新编译回静态图。
但这一次,编译是可选的、自动的、Python 透明的。你不用重写代码、不用学新 API,只要给函数加一个装饰器:
@torch.compile
def model(x):
return torch.relu(x @ W) + b
后面的事,PyTorch 通过 TorchDynamo 拦截 CPython 的字节码、用 AOTAutograd 把正反向图都翻出来、再用 TorchInductor 生成融合后的 Triton kernel —— 既保留了动态图开发的舒适,又拿回了静态图的性能。
理解 torch.compile 这套编译器栈,就是理解今天 PyTorch 团队最聪明的工程师在花心思做什么。这也是为什么本书要在第六篇用整整 4 章来拆解它。
为什么是源码而不是论文
PyTorch 没有官方”白皮书”。它有一篇 NeurIPS 2019 的论文(《PyTorch: An Imperative Style, High-Performance Deep Learning Library》),但那是项目早期的浓缩,描述的是 v1.0 时代的设计。今天的 PyTorch 已经长得几乎是另一套系统:dispatcher 重写过、autograd 引擎重写过、torch.compile 整个生态从无到有 —— 真正承载当代 PyTorch 设计的,只有源码本身。
这是 PyTorch 与传统系统的不同之处:
| 维度 | 传统系统软件 | PyTorch |
|---|---|---|
| 主要参考 | RFC + 论文 + 书 | 源码 + PR 讨论 + Dev Pod 视频 |
| 设计文档 | 完备且权威 | 散落在 issue、design docs、Blog post |
| 核心抽象稳定性 | 几十年不变 | 18-24 个月一次大重构 |
| 学习路径 | 自顶向下 | 自顶向下读源码 + 自底向上看 commit |
这意味着如果你只读教程或者只读论文,你拿到的永远是滞后两年的 PyTorch。要紧跟 PyTorch,必须能自如地读它的源码。
源码读到什么深度算够
读源码不是越深越好。我把”源码深度”分成五个层级:
| 层级 | 描述 | 适合谁 |
|---|---|---|
| L1 | 看官方文档、知道 API 怎么用 | 所有用户 |
| L2 | 知道每个 API 在哪个 Python 文件、调用了哪些子函数 | 进阶算法工程师 |
| L3 | 能跟 Python → C++ → CUDA 的完整调用链、知道 dispatcher / autograd / allocator 这些底层机制 | 系统工程师、本书的核心读者 |
| L4 | 能修改 PyTorch 主仓代码、提交 PR、参与 RFC 讨论 | PyTorch 贡献者 |
| L5 | 能独立维护 PyTorch fork、做自家硬件后端集成 | 大厂底层平台团队 |
本书目标把读者带到 L3。再深的需求超出篇幅,但读完本书之后,你已经具备从 L3 自学到 L4 的所有”地形知识”。
一个真实例子:为什么读源码能救命
举个真实场景:你在用 8 卡 H100 训练一个 70B 的模型,发现训练速度只有理论值的 60%。各种 profile 之后发现 GPU 利用率飘忽不定,平均 70% 左右。
如果只看博客,你能找到的建议大概是 “用 FSDP”、“用 mixed precision”、“调 batch size”、“看看 NCCL 通信”。这些都对,但都是碎片化的猜测。
如果你读过 PyTorch 的源码(特别是本书第 4 章和第 17 章),你会立刻想到几个具体怀疑点:
- CUDA Caching Allocator 在分布式场景下可能在做隐式同步:每次
cudaFree都会触发整设备 sync —— 用expandable_segments=True试试 - DDP 的梯度桶大小默认 25MB,对 70B 模型太小:每个 bucket 触发一次 AllReduce,频次太高 —— 调到 100MB 试试
- NCCL 的
NCCL_NTHREADS默认值在 H100 集群上可能是错的:手动调成512 - AOTAutograd 在某些 op 上可能没生成最优 partition:用
TORCH_COMPILE_DEBUG=1拉日志看 partition 决策
这些都是藏在源码里的工程经验,没有任何博客能一次告诉你这么细。读过源码的人可以从问题描述快速跳到嫌疑代码段,半小时定位的问题,不读源码的人要试一周才能猜到方向。这就是源码的价值。
本书的目标就是把”读源码”这件事从苦力变成乐趣 —— 给你一张地图,标好哪些路是干道、哪些路是死胡同、哪些拐角藏着真正的宝藏。
读源码的”工作流”
一个常见的误区是:以为读源码就是从 __init__.py 开始,一行一行往下翻。这种线性读法对 100 行的小项目可能可行,但对 PyTorch 这种v2.11 实测约 110 万行 C++/Python 混合的代码库就是灾难 —— 你会在第一周耗光精力,连 dispatcher 的入口都没找到。
真正可行的工作流是带着问题读:
- 先从 Python 用户视角写一段最小复现代码,比如 “我想知道
tensor.add是怎么调起来的”,写一个 5 行的 demo - 用
gdb/lldb或py-spy给关键函数下断点,在 Python 层和 C++ 层都打栈 - 从断点回溯调用栈,把每一层的关键文件路径记下来 —— 这就是这个子系统的”骨架”
- 挑骨架上的 1-2 个核心文件精读,其他先 grep 关键词跳过
- 写一段微型代码验证你理解的行为,看预期与实测是否一致
这就是本书每一章背后的方法论。每一章你都会看到:先一段最小复现代码 → 调用栈或时序图 → 几个核心文件精读 → 一个微型实验。章节结构的本身就是教你怎么读源码。
几条避坑提示
读 PyTorch 源码时,有几条坑是大多数人会踩的:
- 不要被 C++ 模板劝退:PyTorch 的 C++ 用了大量 SFINAE、CRTP、模板元编程。但 90% 的核心逻辑读懂宏展开后就是 if/else。第一次见
c10::guts::if_constexpr这类 utility 觉得头大很正常 —— 跳过实现细节、看它做了什么就行 - 不要在生成代码上纠结:PyTorch 的 ATen 算子大量是 代码生成出来的(从
native_functions.yaml)。你在build/目录里能看到生成后的Operators_0.cpp等文件长达数万行 —— 这些是机器写的,没必要逐行读,理解生成规则即可 - 不要小看注释:PyTorch 关键代码的注释非常详细,许多设计决策直接写在头文件里。比如
c10/core/DispatchKey.h顶上有 200 多行注释把 DispatchKey 的设计哲学讲得清清楚楚 —— 比任何博客都权威 - 不要忽视 PR 描述:每一次大重构都会有一份长 PR 描述,写明动机、设计、性能数据。比如 V1 dispatcher 的引入是 PR #34707,AOTAutograd 的进入是 PR #76535 —— 这些 PR 描述就是 PyTorch 的”设计文档”,但散落在 GitHub 上没人系统整理
本书的读者画像
本书面向这几类人:
| 读者 | 典型问题 | 对应章节 |
|---|---|---|
| AI 算法工程师 | ”我的训练为什么慢?为什么 OOM?“ | ch04 / ch11 / ch21 |
| AI 系统工程师 | 训练框架选型、自定义算子、平台对接 | ch05 / ch22 / ch16-ch18 |
| 大模型研究者 | 想理解 FSDP / 混合精度 / 分布式训练原理 | ch15 / ch17-ch20 |
| 编译器爱好者 | torch.compile 的字节码拦截、IR 设计、Triton codegen | ch12-ch15 |
| 推理优化工程师 | 训练侧的内存管理与导出格式如何影响推理 | ch04 / ch06 / ch19 / ch20 |
| 技术总监 / 架构师 | 评估 PyTorch 与 JAX / 国产框架的工程权衡 | ch01 / ch23 |
前置知识
你不需要精通 C++ 模板元编程,但需要具备:
- 熟练 Python(能读 metaclass、descriptor、
__torch_function__这一类协议) - 基本 C++(能看懂模板、
std::shared_ptr、Move 语义、虚函数 —— 但本书会刻意避开 C++ 黑魔法的炫技路径) - 深度学习基础(什么是矩阵乘法、什么是反向传播链式法则、什么是 ReLU)
- 基本的 GPU 知识(什么是显存、什么是 SM、什么是 stream、什么是 kernel 启动)
- 操作系统基础(进程、线程、共享内存、信号 —— DataLoader 多 worker 章会用到)
如果其中某些点你不熟悉,没关系,本书在每次首次出现时都会简短铺垫。但完全没接触过深度学习的话,建议先读一遍 Andrej Karpathy 的《Neural Networks: Zero to Hero》课程,再回来读本书。
不需要的前置
为了避免不必要的劝退,明确说一下你不需要会的:
- 不需要精通 CUDA C++ —— 本书会出现 CUDA 调用,但只关注它怎么被 PyTorch 调度,不需要你能从零写一个 reduction kernel
- 不需要懂 NCCL 内部 —— 第 16 章会讲 PyTorch 怎么集成 NCCL,但不展开 NCCL 的环形拓扑算法(那是另一本书)
- 不需要会 LLVM —— TorchInductor 不直接对接 LLVM,它生成 Triton DSL,由 Triton 编译器去对接 LLVM
- 不需要熟悉 X86/ARM 汇编 —— PyTorch 的优化大量集中在 GPU,CPU 后端有 oneDNN / FBGEMM 这种成熟库支撑,本书只会在必要时简短带过
如果有一天你想再下一层,这些主题各有自己的专著(NVIDIA 的 CUDA 编程指南、《Triton》论文、LLVM 官方教程) —— 本书的边界在 PyTorch 这一层。
与丛书其他书的关联阅读
本书是 AI 训练 + 推理双子卷 的训练侧,与已经成书的 《vLLM 内核探秘》 互补:
- 共享部分(Dispatcher / ATen 算子 / CUDA Allocator / NCCL):在两本书里看到不同的视角 —— vLLM 看到的是”调用方”,本书看到的是”实现方”
- 训练独有(Autograd / DDP / FSDP / torch.compile):本书完整覆盖
- 推理独有(PagedAttention / KV Cache / 投机解码):vLLM 那本完整覆盖
把两本书放在一起读,你会看到一个完整的 LLM 工程世界 —— 训练时千卡跑出权重,推理时千机服务用户,中间靠 PyTorch 的张量、HuggingFace 的格式与 vLLM 的引擎串起来。
更广泛的关联:
- 《Tokio 异步运行时》:autograd Engine 用了和 Tokio 类似的 work-stealing 多线程调度,第 8 章会做对照
- 《Rust 编译器之路》:TorchInductor 的 Lowering 与 Cranelift 的 IR 设计有同源思想,第 14 章会做对照
- 《Serde 元编程》:ATen 的
native_functions.yaml→ C++ 注册的代码生成路径,与 Serde 派生宏哲学相通,第 6 章会做对照 - 《MCP 协议剖析》:分布式训练里的 RPC 概念与 MCP 的传输层在某些设计选择上类似
本书的章节组织
全书 23 章 + 前言,按 “从外到内 → 从基础设施到编译器栈 → 从单机到分布式 → 从核心到周边” 的顺序组织:
graph TD
P[前言]
P --> P1[第一篇 全景<br/>第1章]
P1 --> P2[第二篇 张量层与运行时<br/>第2-4章]
P2 --> P3[第三篇 Dispatcher 与算子<br/>第5-6章]
P3 --> P4[第四篇 Autograd<br/>第7-8章]
P4 --> P5[第五篇 nn 与训练栈<br/>第9-11章]
P5 --> P6[第六篇 torch.compile<br/>第12-15章]
P6 --> P7[第七篇 分布式训练<br/>第16-18章]
P7 --> P8[第八篇 工程与生态<br/>第19-23章]
style P1 fill:#fef3c7
style P2 fill:#dbeafe
style P3 fill:#dbeafe
style P4 fill:#dcfce7
style P5 fill:#dcfce7
style P6 fill:#fce7f3
style P7 fill:#fce7f3
style P8 fill:#f3e8ff
第一篇 全景(第 1 章)
第 1 章先把 PyTorch 的全貌摊开:从一个 c = a + b 的张量加法切入,俯瞰 Python 前端 → Pybind11 桥 → c10 抽象 → Dispatcher → ATen 算子 → CUDA kernel → CUDA Caching Allocator 整条调用链。这一章看完,后面所有章节都能在这张地图上找到位置。
第二篇 张量层与运行时(第 2-4 章)
PyTorch 的”皮肤”是 Tensor,“骨架”是 c10。第 2 章拆 Tensor / Storage / TensorImpl 的三件套结构 —— 为什么 view 操作不复制内存、为什么 .data_ptr() 拿到的指针在多次 view 后能保持一致。第 3 章下到 c10 层:Device、ScalarType、Layout、DispatchKey、intrusive_ptr —— PyTorch 不用 std::shared_ptr 而是自己造了 intrusive_ptr,理由值得细说。第 4 章解剖 CUDA Caching Allocator —— PyTorch 不调原生 cudaMalloc,而是有一个自己的”显存池”,理解它就能理解为什么 torch.cuda.empty_cache() 在大多数时候是无效的。
第三篇 Dispatcher 与算子(第 5-6 章)
PyTorch 的”心脏”是 Dispatcher。一个张量加法被调用时,它必须在十几个候选实现中选一个:CPU 还是 CUDA?float32 还是 bfloat16?要不要走 autograd?要不要走 torch.compile?要不要被 functionalize?这一切由 Dispatcher 用一个 DispatchKeySet 位图 + key 优先级排序的极简机制完成。第 5 章把它讲透。第 6 章则讲 PyTorch 怎么用 native_functions.yaml + 代码生成 一次性生成数千个算子的 C++ 与 Python 包装代码 —— 这是 PyTorch 工程化的最佳实践之一。
第四篇 Autograd 自动求导(第 7-8 章)
loss.backward() 是每个 PyTorch 用户每天都在调用、却很少有人能讲清楚原理的接口。第 7 章讲 反向图怎么在前向时被偷偷地构建 —— 每个支持 autograd 的算子都会在前向时记录一个 grad_fn,这些 grad_fn 串成一张反向图。第 8 章讲反向图怎么被 多线程 work-stealing 引擎 执行 —— 这部分代码在 torch/csrc/autograd/engine.cpp,它的设计与 Tokio 调度器有着惊人的相似。
第五篇 nn 与训练栈(第 9-11 章)
第 9 章讲 nn.Module 这个 PyTorch 用户最熟悉的类是怎么实现的:__setattr__ 偷偷拦截参数赋值、Hooks 怎么挂载、state_dict 怎么递归收集。第 10 章讲 torch.optim 各种优化器(SGD / Adam / AdamW)的实现,以及最近几个版本引入的 Foreach / Fused / Capturable 三种性能模式。第 11 章讲 DataLoader 的多进程 worker —— 它用 torch.multiprocessing 启动 worker、用 _utils.worker.WorkerInfo 同步状态、用 pin_memory_thread 异步搬运数据,是一台精密的多进程机器。
第六篇 torch.compile 编译器栈(第 12-15 章)
这是本书的技术高峰。第 12 章讲 TorchDynamo 怎么用 CPython 3.11+ 的 PEP 523 帧评估 API 拦截 Python 函数调用,逐字节码解析、构建 FX Graph、生成 guards。第 13 章讲 AOTAutograd 怎么把这个图做”函数化”(去除原地操作)、再用 min_cut_rematerialization_partition 切分出最优的正向 / 反向子图。第 14 章讲 TorchInductor 怎么把 ATen IR 一路 lower 到 Triton DSL、再编译成 PTX。第 15 章把整条链路串起来,并讲 CUDA Graph 与 torch.compile 的协作 —— 这是 vLLM 等推理引擎获得 30%+ 加速的关键。
第七篇 分布式训练(第 16-18 章)
第 16 章讲 PyTorch 的分布式底座:进程组、torch.distributed.init_process_group、ProcessGroupNCCL 怎么把 NCCL 集成进 PyTorch 的张量。第 17 章讲 DDP:DistributedDataParallel 怎么用反向传播 hook 触发梯度的环形 AllReduce、用 梯度桶(Gradient Bucket) 把小梯度合并成大消息减少通信开销。第 18 章讲 FSDP:参数怎么被切分到 N 个 rank、前向时怎么 unshard、反向时怎么 reshard、为什么这种”用通信换显存”的策略在大模型训练中是必须的。
第八篇 工程与生态(第 19-23 章)
最后一篇收尾整个训练栈外围。第 19 章讲序列化:从 pickle 到 safetensors,从 torch.save 到 torch.distributed.checkpoint。第 20 章讲量化与混合精度:autocast 上下文管理器、PT2E 量化流程、bf16/fp8/int8 的硬件支持。第 21 章讲 Profiler 与 Kineto:怎么用一行 torch.profiler.profile 拿到 ChromeTrace,怎么读懂 GPU 时间线。第 22 章讲怎么写一个真正”生产级”的自定义算子 —— 用 TORCH_LIBRARY 注册、加 autograd、加 meta 函数、加 functionalize 规则、加 inductor lowering。第 23 章作为收官,盘点 PyTorch 一路走来的设计决策与未来方向。
阅读路径
不同读者推荐不同的阅读顺序:
flowchart LR
A[前言 + ch1] --> B{读者类型}
B -->|算法工程师| AB[ch9 → ch10 → ch7 → ch11 → ch15 → ch20]
B -->|系统工程师| SB[ch1-ch6 → ch22 → ch4 → ch16-18]
B -->|编译器爱好者| CB[ch1 → ch5-6 → ch12-15 → ch23]
B -->|大模型训练| LB[ch1 → ch7-8 → ch16-18 → ch20 → ch15]
style AB fill:#dcfce7
style SB fill:#dbeafe
style CB fill:#fce7f3
style LB fill:#fef3c7
如果你只有一个周末的时间:读前言 + 第 1 章 + 第 5 章(Dispatcher)+ 第 7 章(Autograd)+ 第 23 章。这五章拼起来就是一份”PyTorch 心智模型”的最小闭环 —— 之后再按需要回来读其他章。
PyTorch 的工程文化:理解源码的另一个维度
读 PyTorch 源码,除了看代码本身,还要看代码背后的工程文化 —— 这种文化解释了为什么源码长成现在这样,而不是另一种合理的样子。
极致的向后兼容承诺
PyTorch 团队对 BC(Backward Compatibility)的态度近乎偏执。一个 Python API 一旦发布,几年之内不会被破坏 —— 即便它的 C++ 内部已经被重写过两三次。这种承诺在源码里留下的痕迹是:
- 大量的”兼容层”文件,比如
torch/distributed/fsdp/__init__.py同时导出 FSDP-1 和 FSDP-2 的接口 - 弃用的 API 用
@deprecated装饰器标记并保留(如torch.symeig在 v1.8 弃用,到 v2.x 仍能调用,只是会发警告) - 大量的
if torch.__version__ >= ...条件分支,处理跨版本的迁移
这种”不破坏用户代码”的文化,让 PyTorch 在大公司能放心用 —— 但代价是源码里有大量”历史包袱”。读源码时遇到看似冗余的代码时,先想想”这是不是为某个老 API 留的”,常常能解开困惑。
“性能数据驱动 PR” 的传统
PyTorch 的核心 PR 几乎都附带 benchmark 表格:哪些模型上提速多少、哪些场景下回归。这种”数据先行”的文化让源码演进有迹可循 —— 你能在 PR 描述里看到 “@torch.compile 在 ResNet-50 上提速 1.4x,在 BERT-large 上提速 1.7x,在 Llama-7B 上提速 2.1x,在 batch_size=1 下回归 5%” 这种精确数字。
本书会大量引用这些数据,让”为什么这样设计”的答案有数字支撑,而不是凭感觉。
“小步快跑 + 实验性 flag” 的演进
PyTorch 几乎所有大特性都是先以 torch._inductor / torch._dynamo / torch.compiler / torch._C._functorch 这种带下划线前缀的”内部 API”形式上线,跑一两个版本后再 graduate 到正式 API。比如:
- TorchDynamo 在 1.13 时还叫
torch._dynamo,2.0 起变成torch.compile默认后端 - FSDP-2 在 2.4 时叫
torch.distributed._composable.fsdp,下划线前缀,需要主动 import - DTensor 在 2.0 时叫
torch.distributed._tensor,到 2.8 才去掉下划线
理解这种演进节奏,你就能看懂为什么 PyTorch 主仓里有大量 torch._xxx 模块 —— 它们不是”私有 API”,而是”正在孵化的下一代正式 API”。本书在涉及这些模块时会指出”它将在哪个版本毕业”。
关于 PyTorch 的几个常见误解
写在前面,先把几个广为流传但其实错的说法澄清一下,免得你带着错误前提开始读后面的章节:
误解 1:PyTorch 是”纯 Python”。
完全错。PyTorch 的 C++ 代码量大约是 Python 代码量的 1.5-2 倍。Python 层只是面向用户的接口;真正干活的 dispatcher、autograd 引擎、ATen 算子、内存分配器、CUDA Graph 全部是 C++。Python 端的每一个张量操作最终都会通过 pybind11 跨进 C++。
误解 2:with torch.no_grad(): 就能完全关掉自动求导开销。
torch.no_grad() 关闭了反向图的构建,但 dispatcher 仍然会查 Autograd key(只是命中”什么也不做”的 fallback)。如果你想要进一步关掉这层查找开销,需要用 torch.inference_mode() —— 它把 tensor 标记为 inference 张量,dispatcher 在 key bitmap 上根本就不会进入 autograd 路径。第 5 章会讲这个细节。
误解 3:torch.compile 会让所有代码都变快。
不一定。torch.compile 在以下场景反而可能变慢:函数体很小、调用频次很低(编译开销摊不开)、有大量动态形状(每次都触发重编译)、有 graph break(图被切碎导致每段都需要 launch overhead)。第 12 章会讲怎么诊断 torch.compile 的失效。
误解 4:DDP 比 FSDP 通信少所以更快。
只在小模型上成立。DDP 每一步都要全量 AllReduce 梯度(通信量 = 参数量);FSDP 在前向 / 反向时通过 AllGather + ReduceScatter 来 unshard / 重组(通信量 = 参数量 × 1.5,但每次只传 1/N)。当模型大到单卡放不下时,FSDP 是唯一选项,DDP 根本启动不起来。第 17、18 章会做精确对比。
误解 5:tensor.cuda() 是同步的,需要 torch.cuda.synchronize() 才能确保完成。
tensor.cuda()(即 Tensor.to('cuda'))默认是 异步的 —— 它把 H2D 拷贝排进默认 CUDA stream,立即返回。但后续在同一个 stream 上的 kernel 都会保证看到搬运结果(CUDA stream 内有强一致性)。synchronize() 只在跨 stream 或者你想 CPU 端确认时才需要。这是 PyTorch 异步执行模型最容易被误解的一点。
把这几条记在心里,后面读章节时会少走很多弯路。
字数与篇幅
本书每章正文 8000 字起步,部分核心章节(如 Dispatcher、Autograd Engine、Inductor)超过 12000 字。代码片段、Mermaid 图、Markdown 标记不计入字数。全书 23 章 + 前言估算正文约 25-30 万字,相当于一本中等厚度的技术专著。
我的目标不是写一本”快速入门”,也不是写一本”参考手册” —— 这两类书市面上已经够多了。我的目标是写一本你会反复回来翻的书:第一遍读时建立直觉,第二遍读时核对源码,第三遍读时拿出来当工具书查 ATen 算子注册或 Inductor IR 的细节。
致谢与信息源
本书写作过程中参考的权威信息源,按重要性排列:
- PyTorch 主仓源码 (
github.com/pytorch/pytorchv2.11.0):第一手材料,所有引用都标注精确路径 - PyTorch 官方设计文档(散落在
docs/source/notes/与 GitHub Discussions):dispatcher、autograd、torch.compile 都有官方设计笔记 - 关键 PR 与 RFC:每章会列出重要 PR,作为设计意图的最权威说明
- PyTorch Developer Podcast(Edward Yang 主理,含数百期 PyTorch 内部技术解析)
- PyTorch Conference 历年演讲:Soumith Chintala、Edward Yang、Jason Ansel、Horace He 等核心开发者的 KeyNote
- NeurIPS 2019 PyTorch 论文《PyTorch: An Imperative Style, High-Performance Deep Learning Library》:项目早期设计的浓缩
- Triton 论文与 Triton 官方教程:第 14 章 TorchInductor 章节的下游依赖
- CUDA Programming Guide:底层显存、stream 模型的权威说明
所有引用都标注了出处。如果你在书中看到一段”PyTorch 团队某某说”,那一定有 PR、issue、conference 视频或者博客作为支撑。我对自己的承诺是:绝不编造来源、绝不凭记忆杜撰行号、绝不为流畅而牺牲准确。
如果发现书中有错误(哪怕是行号偏移),欢迎在站点反馈页留言或发邮件。这本书会随 PyTorch 版本跟进维护。
版权声明
本书采用 CC BY-NC 4.0 许可协议。转载或引用请署名 杨艺韬 并附原文链接,禁止商业用途。
“The best way to predict the future is to implement it.”
—— Alan Kay(PyTorch 团队的内部博文常引用这句作为座右铭)
评论 0
还没有评论,来说两句吧。
评论加载失败,刷新重试。