PyTorch：动态图与自动微分

破除误解#

数值微分用差分去逼近导数，步长一选就引入截断误差或舍入误差。符号微分在表达式树上做代数展开，容易把共同子式胀成指数级的纸面公式。反向模式自动微分既不抽样也不展开成巨型闭式：它把每个初等运算的局部雅可比留下，在反向遍历时与上游传来的向量相乘。误差是浮点误差，复杂度与前向计算同阶，而不是与参数个数成正比。把三者混称“求个导”，会看不见 PyTorch 为什么必须维护一张图。

误解二：动态图比静态图“没有编译、所以更慢”，如此而已。

早期 TensorFlow 的静态图先定义后执行，优化器看得见整张图，却让控制流变成图里的特殊节点，Python 调试器也够不着中间值。PyTorch 的 define-by-run 让 if 和 Python 循环就是控制流，磁带只记录这次实际走过的边。代价是：每一次前向都可能长出形状不同的图，能做的跨算子融合默认更少。后来的 torch.compile（TorchDynamo + Inductor）试图在动态性上再抓一份图来编译，那是承认磁带灵活、也承认灵活有税。快慢不是道德问题，是这次前向你愿不愿意把图钉死。

误解三：PyTorch 会替你在集群上排作业、换掉失败的 worker。

不会。DistributedDataParallel 解决的是一次作业内部如何同步梯度，它假设进程组已经存在。GPU 从哪来、进程谁拉起、机器挂了谁重启，是 Kubernetes、Slurm 或训练平台的事。把训练框架误认成集群操作系统，会在第一步就把故障域画错。autograd 引擎只对本地磁带负责。

磁带上到底记了什么#

一次前向里，每个需要梯度的张量带着一个 grad_fn，指向创建它的那个 Function 节点。节点保存反向时还要用的上下文：卷积的输入、softmax 的归一化分母、矩阵乘的另一侧。这些“被保存的张量”往往比参数本身更大。训练时显存爆掉，第一嫌疑经常不是模型权重，而是磁带上的激活。梯度检查点故意丢掉一段激活，反向时再前向算一遍，用计算换内存。这是在承认：磁带的默认策略是“全记住”，而不是“只记最小值”。

反向从标量损失出发，相当于用 1 作为初始向量，沿拓扑逆序做向量—雅可比乘积。若损失不是标量，必须提供 gradient 参数，明确你要的是哪一个线性组合。这不是 API 挑剔，这是反向模式的定义：它高效计算的是 vᵀ J，不是整张雅可比。参数量是百万到十亿时，没有人想把 J 物化出来。前向模式自动微分高效的是 J u，适合输入维度低、输出维度高的情况。深度学习的损失是标量、参数是高维，所以磁带几乎总是反向的。

高阶梯度把同一套机制叠起来。create_graph=True 时，反向本身也在磁带上留下节点，于是可以对梯度再求梯度。二阶优化、元学习需要这个。代价是磁带变长、被保存的张量再翻一倍。默认关掉，是因为大多数训练只需要一阶。

原地操作是磁带与可变张量之间的结构性摩擦。Python 用户习惯 x += 1。磁带若仍指向旧存储，反向就会读到新值。PyTorch 用版本计数来发现这种冲突：张量被原地改写，版本加一；反向时发现版本对不上，就拒绝继续。这比静默算错更吵，也比禁止一切原地更灵活。

分发器与 ATen#

Python 里一次 x + y 不会直接跳进某颗 CUDA 核。中间是分发器：按设备、dtype、布局、是否需要梯度、是否处于自动混合精度，选出真正的内核。这张表的 C++ 层叫 ATen。自动求导键往往先于设备键：它先把节点接到磁带上，再把调用重新分发给真正算数的内核。顺序错了，梯度会丢，或者会在不该建图的推理路径上建图。

算子清单由代码生成维护。一份描述文件生成 Python 绑定、C++ API、梯度包装和分发注册。用户以为 add 是一个函数，底下是一组按键切开的实现。Edward Yang 把这套优先级写成可扩展注册表，第三方设备不必改引擎。自定义算子若只挂了 Python 绑定、没向分发表登记，编译器只能把它当成不透明黑盒，图在这里断开。

autograd 引擎的热路径在 C++ 里维护就绪队列：某个节点的所有上游梯度到齐，就可以执行它的 backward。for 循环里反复创建微小张量会很疼：每一次都是一次 Python 到 C++ 的节点登记。把循环写进一个融合内核，磁带上就只剩一个节点。

Inductor：2020 年代把足迹再抓成图#

PyTorch 2.0（2022 年宣布，Ansel 等人 ASPLOS 2024）把 torch.compile 做成稳定入口。TorchDynamo 挂钩 Python 帧求值，把字节码翻译成 FX 图；碰上它不理解的 Python 就切开，回退到 eager。AOTAutograd 在编译期把前向和反向一起追踪。Inductor 把图降到循环级 IR，GPU 上生成 Triton 融合核，CPU 上生成 C++/OpenMP。一次 layer-norm 在 eager 里可能是一串小内核，在 Inductor 里常常融成一个。

动态图没有消失。它仍是默认交互面，编译器是后加的加速器。图断裂太多，装饰器就变成带额外开销的 eager。动态形状是第二道悬崖：按具体长度编译，会在直方图每个柱上各生成一份代码。科研要随时打印张量；编译器要没有隐藏副作用。作为系统对象，它卖的不是单纯的动态或静态，而是允许你在这条光谱上滑动，并看清每一档的账单。

和 JAX 对照：JAX 把函数当成纯函数，用变换生成新函数，XLA 拿到稳定描述。PyTorch 把图当成这次执行的足迹。调试友好是足迹的收益；整图内存规划默认更弱。两条哲学都诚实，诚实的位置不同。

被保存的张量、检查点与跨卡合同#

训练时显存爆掉，第一嫌疑经常不是模型权重，而是磁带上的激活。一次 Transformer 前向里，注意力的归一化分母、MLP 的中间激活、残差上的层归一化，都会被节点保存下来供反向使用。梯度检查点把某几段标成“只留输入”：反向走到这段时先再算一遍前向，再继续反传。计算量上升，激活占用下降。哪些段该检查点，取决于重新计算的代价和那一段激活的体积。把整网都检查点，训练变成两次前向的慢动作；完全不检查点，长序列上会先撞墙。

原地操作与视图是磁带与可变存储之间的缝。切片、转置、expand 往往不分配新存储。你若对视图做原地写入，版本计数可能拦住你，也可能漏掉，漏掉时反向读到被改写的值。复制一份再改，账单变大，语义变清楚。命令式 API 让这两种选择看起来一样像 Python，运行时却只认存储上的版本号。

数据并行里，autograd 的职责在卡内结束：它把局部 vᵀ J 送到这张卡的叶子，跨卡求和是集合通信，通常挂在 backward 的钩子上。流水线并行把磁带沿层切开，上游向量来自下游阶段而不是本地标量损失。引擎并不理解“这是一张被切开的矩阵”。正确性依赖算子作者把切分后的伴随写对。把“框架会自动处理分布式”理解成引擎定理，会在流水线气泡里找不到梯度。

混合精度改写的是节点上的数值类型。前向用半精度跑矩阵乘，反向常在更宽的类型里累加。缩放损失给磁带上那条标量乘一个系数，反向再除回去。autograd 并不理解“这是 AMP 政策”，它只对实际登记的 Function 做伴随。类型一错，梯度不会举起一面叫混合精度的旗，它只会偏掉。gradcheck 在半精度下更吵：数值微分的步长和舍入叠在一起，抽查要放宽容差，不能把容差当证明。

多次反向是磁带的另一种寿命。默认 backward 释放中间节点，第二次再反传会发现图已被拆。留下图会让显存账单翻倍。训练循环几乎只需要一次一阶梯度。高阶梯度打开 create_graph 时，反向自己也变成磁带上的节点，被保存的张量再翻一倍。默认关掉，是因为大多数人只要一阶，而默认开启会让一次普通训练的显存账单变得无法解释。

叶子、非叶子、detach、no_grad 构成一套故意不对称的 API。推理时关掉磁带，记住激活纯属浪费。微调时只给最后几层 requires_grad=True，磁带只会从那些叶子往回长。自动微分的“自动”，从来都是相对于一组被标记的变量，不是相对于程序里每一个张量。

生产对照：同一损失，两台不同的机器#

eager 路径上，每一次前向都重新登记节点，Python 调试器能停在任意一行，中间张量可以当场打印。编译路径上，同一段 Python 被抓成图，融合会改变内核边界和舍入顺序。损失看起来接近，并不等于每层激活一致。最有效的核对是把同一步同时跑在两条路径上，从输入向损失逐段比较。对不齐的第一处，往往就是图断裂、融合或某次静默的 dtype 提升。把这一步写成持续集成，比争论编译器有没有改语义更接近运维。

数值对比必须锁死分发键。设备、布局、是否需要梯度、是否处于自动混合精度，任一键不同，比较的就不是同一台机器。自定义算子若只在 eager 下有反向、在编译图里变成黑盒，训练会在装饰器打开的那一天突然“不下降”。这不是优化器坏了，这是磁带在那一层断了。科研原型可以停在 eager；要进入稳定训练，必须让那一层同时提供可追踪的前向和伴随，并且向分发表登记。

钩子与通信的顺序同样属于合同。数据并行的 allreduce 必须等到局部梯度写完再启动，也必须在优化器读取 .grad 之前结束。包装模块把这些钩子藏起来，是为了不让用户在磁带和网络之间手写竞态。一旦有人绕开包装直接操作底层参数，或在钩子之外再做一次手动反向，通信与磁带的约定就碎了。分布式没有让链式法则失效，它只是让每个局部雅可比都带着一张切分方案。引擎仍然只做组合。组合的前提是每一段局部伴随都已登记，缺一段就整条链停在那里。集群调度、弹性扩缩、抢占重启，始终在这台机器外面。

dispatcher：同一算子如何变成许多内核#

用户写下的加法不是一个函数，是一扇门。门后按分发键选择实现：设备、数据类型、内存布局、是否需要梯度、是否处于自动混合精度。自动求导键往往先于设备键：它先把节点接到磁带上，再调用下层真正算数的内核。顺序错了，梯度会丢，或者会在不该建图的推理路径上建图。ATen 把这些内核收成与 Python 无关的算子目录；Python 前端只负责把用户代码翻译成对目录的调用。

这张表是 PyTorch 能同时服务科研与芯片厂商的原因，也是启动时间和二进制体积的来源。键一多，注册和兜底逻辑一起长。自定义算子若只挂了 Python 绑定、没向分发表登记模式，编译器只能把它当成不透明黑盒，图在这里断开。生态里大量科研扩展至今停在这个状态：即时模式能跑，编译形同虚设。把「PyTorch 算子」理解成一个 C++ 函数，会看不懂为什么同一句加法在训练、推理、半精度与编译路径上会走进四个不同的实现。

跨域连接#

• 线性代数：反向模式做的每一件事都是一次伴随作用。局部导数是雅可比，上游传来的是向量，引擎输出的是 vᵀ J。广播和缩减之所以难写对，是因为它们对应的伴随不是“把形状改回去”那么直观：沿着被广播的轴要做求和，沿着被缩减的轴要把梯度铺开。卷积、注意力、批归一化的反向公式，全是在写相应线性映射的伴随算子。把 autograd 只看成软件功能，会错过它真正在实现的那条代数恒等式。
• 证明：链式法则是一条定理，磁带是这条定理的一次算法化。定理要求每个局部映射可微，并且组合按固定顺序发生。Python 控制流让顺序在运行时才确定，于是“证明”只对这一次足迹成立：下一次若走了另一分支，那是另一条组合，要另建一张图。gradcheck 用数值抽样充当廉价的反例搜索，不是对可微性的证明。自定义 backward 一旦写错，链式法则在数学上仍真，在这台机器上已经假了。
• 范畴论：可微映射在组合下封闭，这正是磁带敢于按节点拼接的原因。每个 Function 是一枚箭头，backward 是它的伴随箭头；引擎做的是沿相反范畴把上游余向量拉回来。动态图等于每次前向现场选择一条可复合的路径，而不是事先给出对象之间的全部箭头。create_graph=True 则要求伴随箭头自己也还在那个可微范畴里，而不是掉进一次性的数值过程。范畴语言在这里标定了接口：你能接入引擎，当且仅当你同时提供正向箭头和伴随箭头。
• 信息几何：训练循环把损失看成参数流形上的标量场，autograd 给出的是这个场在当前点的坐标梯度。自然梯度还要把这组坐标梯度用 Fisher 信息改度量，那一步已经超出默认磁带。但磁带决定了你能否在同一前向里取出对数似然对参数的导数，从而估计那项度量。信息几何关心的“用哪套度量来走一步”，和 autograd 提供的“先把坐标梯度算出来”，是相邻的两层。没有后者，前者只是流形上的一句宣言。

参考文献#

• Paszke, A. et al. Automatic Differentiation in PyTorch. NeurIPS Autodiff Workshop, 2017.
• Paszke, A. et al. PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS, 2019.
• Griewank, A. & Walther, A. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation. SIAM, 2nd ed., 2008.
• Baydin, A. G., Pearlmutter, B. A., Radul, A. A., Siskind, J. M. Automatic Differentiation in Machine Learning: A Survey. Journal of Machine Learning Research 18, 2018.
• Ansel, J. et al. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. ACM ASPLOS, 2024.

延伸阅读#

• Bradbury, J. et al. JAX: Composable Transformations of Python+NumPy Programs. 2018 起开源文档。（对照：纯函数变换 vs 磁带）
• PyTorch 官方 Autograd mechanics 文档（torch.autograd 的原地操作、保存张量与多次反向规则）.
• Edward Yang, Let's talk about the PyTorch dispatcher（公开技术写作, 2020）。分发键与再分发的实现口径。
• Rumelhart, D. E., Hinton, G. E., Williams, R. J. Learning Representations by Back-propagating Errors. Nature 323, 1986.（反向传播作为反向模式 AD 的特例）