收尾 · 从 Python 到金属，整台引擎的全景

收尾 · 从 Python 到金属，整台引擎的全景

漫游到头了。我们从最底下的一块内存，一路爬到最上面的 torch.compile。这一章把沿途看到的齿轮，拼成一张完整的图。

一张图：PyTorch 的全栈

从你敲的 Python，到 GPU 上的金属，PyTorch 大致是这样层层叠起来的：

每一层，都是这趟漫游的一章：

• Tensor 与 Storage —— 地基：Tensor 是「怎么读一块扁平内存」的说明书(sizes/strides/offset),view/transpose 零拷贝、contiguous 的由来全在这。
• dispatcher —— 中枢：按分发键把一次调用路由到正确内核；autograd、autocast 只是链上的拦截层。
• ATen 与内核 —— 车间：两千多个算子靠 native_functions.yaml 声明 + codegen 生成 + TensorIterator 复用。
• autograd —— 引擎：前向悄悄建图，反向按拓扑逆序回放、复合局部 VJP。
• nn.Module —— 上层糖：靠 __setattr__ 自动登记参数/子模块/buffer，自己不碰数学、把求导全甩给 autograd。
• CUDA 与显存 —— 金属：异步入队的流、以及那个「囤内存复用」的缓存分配器。
• torch.compile —— 加速层：架在 eager 引擎之上，能编译就编译、不能就落回地基。

追一行 loss.backward() 穿过全栈

把这些层用一次真实调用串起来，你就真的「看穿」了 PyTorch。以 nanoGPT 的一步训练为例：

1. logits, loss = model(X, Y) —— nn.Module.__call__ 跑 forward；里面每个张量运算，都经 dispatcher 路由，先被 Autograd 键拦下记一笔账、建图，再 redispatch 到 CUDA 内核在 GPU 上异步执行。前向走完，一张挂满 saved 激活的反向图，连好了。
2. loss.backward() —— autograd 引擎从 loss 播下种子梯度 1，逆着图按拓扑序回放：每个节点用链式法则算 VJP、把梯度往上游传、多路相加，最后 AccumulateGrad 堆进各参数的 .grad。图随即释放，显存回收。
3. optimizer.step() —— 优化器拿着 .grad 更新参数（Parameter，也就是图的叶子）。
4. optimizer.zero_grad() —— 清空 .grad，因为它是累加的。

nanoGPT 训练循环那几行，你当时读的是「上层剧本」；现在，你看清了幕后每一层的运作。

你带走了什么

漫游的收获，不是背下多少文件路径，而是一套「往下看」的心智模型。以后再碰到这些，你不再抓瞎：

• autograd 报错（grad_fn、AccumulateGrad、backward a second time）—— 你知道图是什么、什么时候建、什么时候释放。
• CUDA OOM —— 你知道显存被参数、优化器状态、激活、分配器囤块四坨吃掉，也知道每个省显存技巧在拧哪一坨。
• 给 GPU 代码计时不准 —— 你知道要 synchronize，因为调用是异步入队的。
• 想加自定义算子 / 反向 —— 你知道从 torch.autograd.Function 或 native_functions.yaml 那一层介入。
• torch.compile 不提速 —— 你知道去查 graph break。

PyTorch 从一个「大概会用」的黑箱，变成了一台你看得见齿轮、能对着某一层去调试和优化的机器。

三部曲，合成一张完整的地图

到这里，三个 AI 源码专栏凑齐了，它们本就是一套：

• nanoGPT 源码解读 —— 一个真实的模型，六百行，逐行读懂 GPT 的本体。
• AI 大模型解构 —— 真实大模型在这副骨架上改了什么（MLA、MoE、归一化……），给你一把对比的尺子。
• PyTorch 源码漫游 —— 托起这一切的引擎，loss.backward() 底下到底怎么转。

模型、模型的变体、以及跑模型的引擎——三层都看穿了。 从此，「大模型」对你而言，不再有哪一层是纯粹的黑箱：上面是可逐行读的代码，中间是可逐条追问的设计选择，底下是可看清齿轮的引擎。

最后一句

PyTorch 那件 Python 外衣太好穿，好穿到你几乎忘了底下还有一台庞大的引擎。这趟漫游想留给你的，就是一种踏实：那台引擎不神秘。它有清晰的分层、可追的源码、可讲的道理。 需要的时候，你随时能掀开外衣，下去看看齿轮怎么咬合。

《PyTorch 源码漫游》全篇完。 也感谢你读完这套三部曲——愿你从此看 AI，多了一层「看得见底下」的从容。

参考文献

• [1] Edward Z. Yang. “PyTorch internals.” 2019. http://blog.ezyang.com/2019/05/pytorch-internals/
• [2] Paszke et al. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” NeurIPS 2019. arXiv:1912.01703.
• [3] Ansel et al. “PyTorch 2: ... Dynamic Python Bytecode Transformation and Graph Compilation.” ASPLOS 2024. DOI: 10.1145/3620665.3640366.
• 各章文末附有该章直接引用的文献与源码位置。