第7章 Autograd 设计原理与反向图构建

"When two or more Edges (from different sources) point at the same input to a Node, the values produced along all of these edges are implicitly summed prior to being forwarded to the target Node."

—— torch/csrc/autograd/function.h:78-80（v2.11.0）

7.1 一行代码引发的迷思

每个 PyTorch 用户都写过这样的代码：

import torch

a = torch.randn(2, 3, requires_grad=True)
b = torch.randn(3, 2, requires_grad=True)
c = a @ b
loss = c.sum()
loss.backward()

print(a.grad)   # 神奇地有了！
print(b.grad)   # 也有了！

loss.backward() 这一行做了什么？严格地讲：

1. 它怎么知道反向应该从 loss 开始？
2. 它怎么知道 loss 是从 c.sum() 来的？c 是从 a @ b 来的？
3. 它怎么算梯度？matmul 的反向公式从哪里来的？
4. 它怎么知道 把 a 的梯度填到 a.grad？

答案是 PyTorch 在前向时偷偷建了一张反向图。这张图记录了"loss 是怎么从 a、b 算出来的"，反向时沿图反走，应用每个算子的反向公式。本章把这张图的构建过程彻底拆开。

graph BT
    subgraph Forward["前向 (a @ b → c.sum() → loss)"]
        F1[a 张量<br/>requires_grad=True]
        F2[b 张量<br/>requires_grad=True]
        F3[c = a @ b]
        F4[loss = c.sum]
        F1 --> F3
        F2 --> F3
        F3 --> F4
    end
    subgraph Backward["反向图 (在前向时偷偷建好)"]
        B0[GraphRoot<br/>backward 起点]
        B1[SumBackward0]
        B2[MmBackward0]
        B3[AccumulateGrad<br/>对应 a]
        B4[AccumulateGrad<br/>对应 b]
        B0 --> B1
        B1 --> B2
        B2 --> B3
        B2 --> B4
    end
    F4 -.grad_fn.-> B1
    F3 -.grad_fn.-> B2
    F1 -.grad_acc.-> B3
    F2 -.grad_acc.-> B4
    style B0 fill:#fef3c7,stroke:#f59e0b
    style B3 fill:#dcfce7,stroke:#22c55e
    style B4 fill:#dcfce7,stroke:#22c55e

7.2 AutogradMeta：张量身上的"反向身份证"

回到第 2 章 §2.4：TensorImpl 有一个 autograd_meta_ 字段：

// c10/core/TensorImpl.h:2920
std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_ = nullptr;

它是 unique_ptr 默认 nullptr —— 不需要梯度的张量通常不分配 autograd_meta。需要时它指向 torch::autograd::AutogradMeta：

// torch/csrc/autograd/variable.h:227
struct TORCH_API AutogradMeta : public c10::AutogradMetaInterface {
    std::string name_;

    Variable grad_;                          // 累积的梯度
    std::shared_ptr<Node> grad_fn_;          // 产生这个张量的 Node (中间张量)
    std::weak_ptr<Node> grad_accumulator_;   // 累积梯度的 Node (叶子张量)

    mutable std::shared_ptr<ForwardGrad> fw_grad_;   // forward-mode 切线
    std::vector<std::unique_ptr<FunctionPreHook>> hooks_;
    ...
};

——这就是张量的"反向身份证"。理解几个字段：

• grad_fn_：如果张量是某个算子的输出（interior），这个字段指向产生它的 Node。a @ b 产生的 c，c.grad_fn 就是 MmBackward0
• grad_accumulator_：如果张量是用户创建的叶子（leaf），这个字段指向一个 AccumulateGrad 节点 —— 反向时所有汇入的梯度都加到 grad_ 上
• grad_：累积的梯度。loss.backward() 跑完后用户读 a.grad 就是读这个字段
• fw_grad_：forward-mode 切线（jvp 用）

grad_fn_ 和 grad_accumulator_ 是互斥的：一个张量要么是某算子的输出（有 grad_fn）、要么是叶子（有 grad_accumulator），不会同时有。

为什么用 weak_ptr 持有 grad_accumulator？因为 leaf 张量本身可能很多（如所有模型参数），而 AccumulateGrad 节点的存活由反向图的引用关系控制。AccumulateGrad 自己强引用着它对应的叶子张量（torch/csrc/autograd/variable.cpp 的 grad_accumulator() 用这个张量构造它），如果叶子张量再强引用 AccumulateGrad —— 两者互相钉住，形成长寿循环。用 weak_ptr 让 AccumulateGrad 在反向图被释放时自动消失，下次反向再 lazy 创建。这种"按需复活"是 PyTorch 内存管理上又一个细节优化。

7.2.1 lazy materialize_autograd_meta

AutogradMeta 不是张量构造时就分配，而是 lazy 分配。打开 torch/csrc/autograd/variable.h：

TORCH_API AutogradMeta* materialize_autograd_meta(const at::TensorBase& self);

（声明在 torch/csrc/autograd/variable.h:120）只有第一次需要写 autograd 信息（如设置 grad_fn、注册 hook）时才分配。这种 lazy 模式让 inference 路径上的张量完全不付 autograd 元数据的代价 —— TensorImpl 里始终只有一个 8 字节的 unique_ptr，堆上那块 AutogradMeta（grad_ / grad_fn_ / grad_accumulator_ / fw_grad_ / 几组 hook 容器）在纯推理路径上根本不会被分配。注意省下的是 host 侧元数据，不是显存。

7.2.2 leaf vs interior：autograd 视角下的两类张量

a = torch.randn(3, requires_grad=True)   # leaf — 用户创建
b = torch.randn(3, requires_grad=True)   # leaf
c = a + b                                 # interior — 算子结果
d = c * 2                                 # interior

print(a.is_leaf)        # True
print(c.is_leaf)        # False
print(a.grad_fn)        # None
print(c.grad_fn)        # <AddBackward0 object>

leaf 张量有梯度累积器（grad_accumulator）但没有 grad_fn。interior 张量反过来。这种区分是反向传播终止条件的体现 —— 反向沿图走到 AccumulateGrad 节点就把梯度填到张量的 grad_ 字段，停下来。

7.3 Node 与 Edge：反向图的骨架

打开 torch/csrc/autograd/function.h:113：

struct TORCH_API Node : std::enable_shared_from_this<Node> {
 protected:   // 节选：以下成员在源码中都位于 protected 段
    virtual variable_list apply(variable_list&& inputs) = 0;

    // 指向"上游" Node 的边
    edge_list next_edges_;

    uint64_t sequence_nr_;         // 创建顺序，反向时按 priority 排序
    uint64_t topological_nr_ = 0;  // 到叶子的最长路径长度，用于剪枝
    ...
};

每个 Node 是一个反向函数对象，apply 方法接收"输出梯度"返回"输入梯度"。next_edges_ 是反向图的核心：

// torch/csrc/autograd/edge.h:14
struct Edge {
    // 省略构造函数、is_valid() 与比较运算符
    std::shared_ptr<Node> function;   // 指向上游 Node
    uint32_t input_nr;                 // 是上游 Node 的第几个输入
};

Edge = (Node 指针, input 索引)。一个 Node 的 next_edges 是它"上游"Nodes 的列表，每条边告诉反向引擎：当我把梯度往上传时，应该传给哪个 Node 的第几个输入。

graph BT
    subgraph CN["MmBackward0 (c = a @ b 的反向 Node)"]
        N["MmBackward0<br/>──────────<br/>apply(grad) {<br/>  return [<br/>    grad @ b.T,  ← input 0<br/>    a.T @ grad   ← input 1<br/>  ]<br/>}<br/>──────────<br/>next_edges_:<br/>  [(AccumulateGrad_a, 0),<br/>   (AccumulateGrad_b, 0)]"]
    end
    AGA[AccumulateGrad_a]
    AGB[AccumulateGrad_b]
    N -- "next_edges[0] (input_nr=0)" --> AGA
    N -- "next_edges[1] (input_nr=0)" --> AGB
    style N fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

注意 input_nr 字段的作用：当多个 Node 的输出共同流到同一个上游 Node 的多个输入时，每条 Edge 必须知道"我对应的是哪个输入位"。否则梯度会错乱。

7.3.1 反向图是 DAG，不是树

考虑 c = a + a：

graph BT
    AB[AddBackward0]
    AGA[AccumulateGrad_a]
    AB -->|"next_edges[0] (input_nr=0)"| AGA
    AB -->|"next_edges[1] (input_nr=0)"| AGA

AddBackward0 的两条 next_edges 都指向同一个 AccumulateGrad_a —— 因为两个输入都是 a。反向时，AddBackward0 计算出 grad_input0 和 grad_input1，两者会被加起来累积到 a.grad。这种"多边汇入同一节点"是 DAG 的标准结构。

torch/csrc/autograd/function.h:78-80 的注释明确说了这件事：

When two or more Edges (from different sources) point at the same input to a Node, the values produced along all of these edges are implicitly summed prior to being forwarded to the target Node.

反向引擎自动处理这种隐式求和，用户代码完全感受不到。

这一行隐式求和语义是反向传播链式法则在多路径场景下的精确实现 —— 数学上 dL/dx = ∑ (dL/dy_i) (dy_i/dx)，引擎用"汇入同一节点的多条边自动相加"实现这条公式。理解这条规则，你就能预测任意 DAG 的反向行为，不需要逐个手动求导。

7.3.2 sequence_nr 与 topological_nr

每个 Node 还有两个看似冗余的字段：

• sequence_nr_：thread-local 单调递增的创建顺序号。反向时按它做优先级排序（让较新创建的 Node 优先执行，启发式地把"靠近 loss"的反向先做完）
• topological_nr_：拓扑序号。按 function.h 的 NOTE [ Topological Number ] 定义，它是"从本节点出发到任意叶子节点的最长路径长度"，叶子（AccumulateGrad）为 0。它的性质是：若存在 X→Y 的有向路径，则 topo_nr(X) > topo_nr(Y)。反过来不成立 —— 所以它只能在 O(1) 内证明"两点之间不存在路径"，用于 autograd discovery 阶段剪掉不可能到达目标的分支

这两个号让反向引擎在保持"DAG 拓扑遍历"语义的同时省下大量无谓遍历。注释里还专门提醒：一个节点一旦被当作别人的 next_edge（有了 parent），它的 topological_nr_ 就不能再变了，源码用 has_parent_ 字段来强制这条不变量。

7.3.3 反向图的"惰性构建"特性

值得强调一个反直觉的事实：前向时建立的反向图，在 loss.backward() 调用前都不真正消耗计算。它就是一些 shared_ptr<Node> 互相指向的对象图。每个 Node 持有几个张量引用（SavedVariable 里）和几条 Edge，但不做任何数学计算。

这意味着即便你前向跑了几百层 transformer 创建了几千个反向 Node，在 loss.backward() 调用前多付的只是这些对象的分配与连接，和前向本身的计算相比通常很小。真正的 backward 计算延迟到 backward 时才发生。

7.4 反向图怎么在前向时被偷偷建

sequenceDiagram
    participant U as 用户代码
    participant V as VariableType wrapper
    participant N as 新 Node MmBackward0
    participant K as ATen kernel mm
    participant T as 输出 Tensor
    U->>V: c = a @ b 触发 mm
    V->>V: 1. collect_next_edges: 收集 a/b 的 grad_fn
    V->>N: 2. new MmBackward0
    V->>N: 3. set_next_edges 连到 a.grad_fn b.grad_fn
    V->>N: 4. SavedVariable 保存 a/b 给反向用
    V->>K: 5. redispatch 到真实 mm
    K-->>V: 输出 tensor
    V->>T: 6. set_history: c.grad_fn = MmBackward0
    V-->>U: c
    Note over T,N: 反向图已经悄悄建好<br/>c.backward 时沿 grad_fn 反走

到了核心问题：前向算子调用怎么变出反向图？答案在第 5 章 §5.5 提过的 VariableType 包装。每个可微算子在 tools/autograd/derivatives.yaml 声明反向规则，由 tools/autograd/gen_variable_type.py 自动生成包装函数。

简化的 VariableType::mm 长这样：

// 简化版生成代码
Tensor mm(const Tensor& self, const Tensor& mat2) {
    // 1. 收集 next_edges (上游 Node)
    edge_list next_edges = collect_next_edges(self, mat2);

    // 2. 创建反向 Node
    auto grad_fn = std::shared_ptr<MmBackward0>(
        new MmBackward0(),
        deleteNode);
    grad_fn->set_next_edges(std::move(next_edges));

    // 3. 保存反向需要的张量
    grad_fn->self_ = SavedVariable(self, false);
    grad_fn->mat2_ = SavedVariable(mat2, false);

    // 4. redispatch 到下一层 (实际计算)
    auto result = at::redispatch::mm(after_autograd_keyset, self, mat2);

    // 5. 把 grad_fn 挂到输出
    set_history(result, grad_fn);

    return result;
}

每一步都要看清楚：

7.4.1 collect_next_edges：连接上游

collect_next_edges 遍历输入张量，每个输入占一个槽位、贡献一条 Edge（torch/csrc/autograd/variable.cpp 的 gradient_edge()）：

• 如果输入是 interior（有 grad_fn_）：Edge 指向那个 grad_fn
• 如果输入是 requires_grad=True 的 leaf：Edge 指向 AccumulateGrad
• 如果输入不需要梯度：grad_accumulator() 返回 nullptr，这个槽位是一条无效 Edge（function 为空，Python 里显示成 (None, 0)）

这就是反向图边连接的全部魔法。前向时每个算子各自决定自己的 next_edges，整张反向图就在不知不觉中拼好了。

为什么这种设计如此简洁？因为反向图的边只取决于"哪个张量是这次 op 的输入"这一信息，而这是每个算子在调用时一定知道的。不需要全局状态、不需要事后分析、不需要 trace pass —— 每个算子各自管自己的 next_edges，整张图自然涌现。

7.4.2 set_history：把 grad_fn 钉在输出上

set_history(result, grad_fn)（torch/csrc/autograd/functions/utils.h）先调 grad_fn->add_input_metadata(result) 拿到输出序号，再调 impl::set_gradient_edge(result, {grad_fn, output_nr})，后者大致等价于：

// torch/csrc/autograd/variable.cpp 的 set_gradient_edge()，节选
auto* meta = materialize_autograd_meta(self);
meta->grad_fn_ = std::move(edge.function);
meta->output_nr_ = edge.input_nr;

这里并不调 set_requires_grad(true)：AutogradMeta::requires_grad() 返回的是 requires_grad_ || grad_fn_，requires_grad_ 只对叶子有意义，interior 张量有了 grad_fn 就自然 requires_grad == True。这一步让 result 也变成 autograd 体系的一员 —— 它有 grad_fn、因而 requires_grad，下一个算子调用时能从 result 推出 next_edges。整个反向图就这样"链式自传染"。

7.4.3 "如果输入都不 requires_grad，跳过整个反向构建"

VariableType::mm 真实代码（生成在 torch/csrc/autograd/generated/VariableType_3.cpp：gen_variable_type.py 按算子名 mm 的 SHA-1 对 5 取模分片，结果是 3）的第一步其实是检查：

auto _any_requires_grad = compute_requires_grad(self, mat2);
std::shared_ptr<MmBackward0> grad_fn;
if (_any_requires_grad) {
    grad_fn = std::shared_ptr<MmBackward0>(new MmBackward0(), deleteNode);
    // ... save tensors, set next_edges ...
}
auto result = at::redispatch::mm(...);
if (grad_fn) {
    set_history(result, grad_fn);
}
return result;

如果所有输入都不 requires_grad，跳过 grad_fn 创建。这种短路让"参数张量混着 inference 张量"的常见 inference 路径不为 autograd 付任何额外开销。理解这条路径，你能精确推理：当模型 weight 关闭 grad、只跑 inference 时，autograd 的 dispatcher 中间层照样经过 VariableType 包装，只是走这条短路 —— 不建 Node、不存张量，但多出来的那一跳 redispatch 还在（§7.7 的表格）。

7.4.4 反向图随调用栈生长，不随时间增长

很多新手会担心"长跑训练的反向图会不会越来越大"。答案是不会 —— 反向图的大小由当前 forward 的复杂度决定，每次 loss.backward() 跑完后各 Node 保存的中间张量随即被释放；Node 对象本身在 loss 等输出张量被下一轮重新赋值、引用归零时整体析构。下一个 batch 的 forward 重新建一张全新的反向图。

例外是 retain_graph=True：用户显式说"我还要再 backward 一次"，反向图保留下来。这种用法在 second-order derivatives、GAN 训练等场景出现，但绝大多数训练不需要。第 8 章 §8.7.1 会讲 retain_graph 在 Engine 里的作用——它只决定 Node 跑完后是否释放 SavedVariable，不影响调度。

7.4.5 跟踪一个真实的 forward → backward graph

让我们用一段最简单的代码完整跟踪反向图构建：

import torch

a = torch.tensor([1., 2.], requires_grad=True)   # leaf
b = torch.tensor([3., 4.], requires_grad=True)   # leaf

c = a + b                                         # → AddBackward0
d = c * 2                                         # → MulBackward0
e = d.sum()                                       # → SumBackward0
e.backward()

每行做的事：

1. a = torch.tensor(..., requires_grad=True)：构造 leaf 张量。AutogradMeta 被分配，grad_fn_=None，grad_accumulator_ 是 weak_ptr → 待 lazy 创建的 AccumulateGrad_a
2. b = ...：同理，对应 AccumulateGrad_b
3. c = a + b：dispatcher 命中 VariableType::add →
   • compute_requires_grad(a, b) 返回 true
   • 创建 AddBackward0 Node
   • collect_next_edges(a, b) 返回 [Edge(AccumulateGrad_a, 0), Edge(AccumulateGrad_b, 0)]
   • grad_fn->set_next_edges(next_edges)
   • redispatch 调真 add，得到 c
   • set_history(c, AddBackward0) —— c.grad_fn = AddBackward0
4. d = c * 2：类似地 → MulBackward0，next_edges = [Edge(AddBackward0, 0), 无效 Edge]（c 是 interior，grad_fn 就是上游 Node）。常数 2 被包成一个不需要梯度的张量传给 mul.Tensor，它仍占第二个槽位，只是那条 Edge 为空 —— d.grad_fn.next_functions 实跑输出 ((<AddBackward0 ...>, 0), (None, 0))
5. e = d.sum()：SumBackward0，next_edges = [Edge(MulBackward0, 0)]
6. e.backward()：从 e.grad_fn (SumBackward0) 出发反向遍历，最终梯度填到 a.grad / b.grad

整张反向图：

SumBackward0  →  MulBackward0  →  AddBackward0  →  AccumulateGrad_a
                                          ↓
                                       AccumulateGrad_b

每个节点的 next_edges 完全在前向时被正确填好。用户写的就是 4 行 Python 代码，PyTorch 在背后偷偷拼出了完整的反向 DAG。

7.5 SavedVariable：前向中间值穿越到反向

很多反向计算需要前向时的中间值。mm 的反向需要 mat2 才能算 grad_self = grad @ mat2.T。但 mat2 是普通 Python 变量，怎么保证它在反向时还活着？

答案是 SavedVariable（类定义在 torch/csrc/autograd/saved_variable.h:22）：

class TORCH_API SavedVariable {
 private:
    at::Tensor data_;                          // 保存的张量数据
    std::shared_ptr<ForwardGrad> fw_grad_;     // forward AD 切线
    std::weak_ptr<Node> weak_grad_fn_;          // 防循环引用 (inplace view 场景)
    uint32_t saved_version_ = 0;                // 保存时的 version_counter
    uint32_t output_nr_ = 0;
    bool was_default_constructed_ = true;
    bool is_inplace_on_view_ = false;
    bool saved_original_ = false;
    bool is_leaf_ = false;
    bool is_output_ = false;
    std::unique_ptr<SavedVariableHooks> hooks_;
    ...
};

它的关键责任：

7.5.1 防 inplace 污染：saved_version_

第 2 章 §2.15 讲过的"version_counter"在这里登场。SavedVariable 在保存时记下 saved_version_：

// torch/csrc/autograd/saved_variable.cpp:48，构造函数节选
was_default_constructed_ = false;
saved_version_ = variable._version();
is_leaf_ = variable.is_leaf();
is_output_ = is_output;
is_inplace_on_view_ = is_inplace_on_view;

反向 unpack 时检查（注意：挂了 saved tensor hooks 就跳过这道检查，因为 hook 之后没法再追踪版本）：

// torch/csrc/autograd/saved_variable.cpp:167，unpack() 节选
if (!hooks_) {
    auto current_version = impl::version_counter(data_).current_version();
    if (saved_version_ != current_version) {
        std::stringstream message;
        message << "one of the variables needed for gradient computation has been "
                   "modified by an inplace operation: [" << data_.toString() << ' '
                << data_.sizes() << ']';
        // ... 再拼上 "which is output N of XxxBackward"、当前/期望 version、
        //     以及"打开 anomaly detection"的提示
        TORCH_CHECK(false, message.str());
    }
}

——这就是那个所有 PyTorch 训练老手都见过的经典报错的来源！它防止用户对 mat2 做 inplace 修改后反向算出错误梯度。比如 b = a.exp(); b.add_(1); b.sum().backward()（ExpBackward0 保存了输出 b），在 CPU 版 torch 2.11.0 上实跑得到：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [3]], which is output 0 of ExpBackward0, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True, check_nan=False).

7.5.2 防循环引用：weak_grad_fn_

考虑 inplace view 场景：

x = torch.randn(3, requires_grad=True)
z = x.clone()                     # 非叶子；对需要梯度的叶子的 view 做 inplace 会直接报错
y = z[:2]                         # view
y.exp_()                          # inplace 改 z 的一部分
z.sum().backward()

y.exp_() 触发的反向 Node（ExpBackward0，公式 grad * result.conj()）需要保存输出 result，也就是 y 自己。而 y 的 grad_fn 经 rebase_history() 又会连回这个 Node —— 强引用会形成循环：ExpBackward0 → SavedVariable(y) → AutogradMeta(y).grad_fn_ → ExpBackward0。

SavedVariable 对这类循环有两手处理，注意区分：

• 一般情况（保存的张量正好是本 Node 的输出、且不是叶子）：构造函数把 saved_original_ 置为 false，只保存 variable.tensor_data()（一份不带 autograd_meta 的浅拷贝）并调 save_metadata() 单独记下元信息。反向时必须由调用方把 grad_fn 通过 unpack(saved_for) 传进来 —— 张量这一侧压根不持有 Node，循环自然不成立
• inplace-on-view 情况（构造时 is_inplace_on_view = true，生成代码传的是 self.is_view()，上例就是这种）：这时才用上 weak_grad_fn_，构造函数里 weak_grad_fn_ = variable.grad_fn();，反向 unpack 时 weak_grad_fn_.lock() 拿回来。源码注释写得很直白："Weak version of grad_fn_ that prevents leaks in rebase_history() for inplace views"

7.5.3 saved variable hooks：接管中间值的存取

PyTorch 还提供 SavedVariable::register_hooks 接口（torch/csrc/autograd/saved_variable.h:49）。

用户可以为每个保存的张量挂自定义 pack / unpack 钩子，Python 侧的入口是 torch.autograd.graph.saved_tensors_hooks。

下面是读者示例，不是 PyTorch 源码原文——pack_to_cpu / unpack_from_cpu 都是示例自取的钩子名：

# 把保存的中间值放到 CPU 而非 GPU，节省显存
def pack_to_cpu(t):
    return t.cpu()
def unpack_from_cpu(t):
    return t.cuda()

with torch.autograd.graph.saved_tensors_hooks(pack_to_cpu, unpack_from_cpu):
    out = model(x)
    loss = out.sum()
loss.backward()

这是大模型训练显存优化的最重要工具之一：把中间值从 GPU 换到 CPU 甚至磁盘，反向时再换回。activation checkpointing（梯度检查点）的非重入实现（use_reentrant=False）就是基于这套 hook 实现的（torch/utils/checkpoint.py 的 _checkpoint_hook 继承自 saved_tensors_hooks）。第 18 章 FSDP 会再用到。

具体的钩子语义：pack_to_cpu 在 SavedVariable 构造时被调用，用户决定怎么"压缩"待保存的张量；unpack_from_cpu 在 backward 用到这个张量时被调用，把"压缩"格式还原。PyTorch 不强制 pack 一定是 device-to-host 拷贝 —— 你完全可以做更激进的事，比如把张量量化成 int8 存下来再反向时反量化（牺牲精度换显存），或者写到 NVMe 磁盘（牺牲速度换更大显存）。PyTorch 自带的 torch.autograd.graph.save_on_cpu 就是这套 hook 的现成封装（它继承自 saved_tensors_hooks，定义在 torch/autograd/graph.py）。

7.5.4 一个真实优化案例：activation_checkpoint

torch.utils.checkpoint.checkpoint(fn, *args) 有两套实现：use_reentrant=False 走上一节的 saved tensor hooks；use_reentrant=True（v2.11 不传该参数时仍回落到它，并发出警告）走 torch/utils/checkpoint.py 里的 CheckpointFunction，一个 autograd.Function。后者的"no_grad 前向 + 反向时重算"思路可以简化成下面这样（示意，不是源码原文）：

class CheckpointFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, fn, *args):
        ctx.fn = fn
        ctx.save_for_backward(*args)
        with torch.no_grad():            # 前向不建反向图
            return fn(*args)

    @staticmethod
    def backward(ctx, *grad_outputs):
        args = ctx.saved_tensors
        with torch.enable_grad():
            outputs = ctx.fn(*args)      # 反向时重新前向，这次建图
        grads = torch.autograd.grad(outputs, args, grad_outputs)
        # backward 的返回值个数必须和 forward 的入参个数一致：
        # forward 的第一个参数是 fn 本身（不可微），要补一个 None
        return (None, *grads)

实际 PyTorch 的 checkpoint 实现更复杂（处理 RNG state、AMP 上下文、use_reentrant=False 的新实现等），但核心就是这套思想。前向时不保存中间值，反向时重算一遍换得显存峰值大幅下降，代价是多跑一遍前向的算力。省多少、慢多少完全取决于你把哪些区段包进 checkpoint，需要在自己的模型上实测。

7.5.5 SavedVariable 的内存账

每一个 SavedVariable 持有 at::Tensor data_（一个引用计数指针，约 8 字节），但它强引用底层 storage。这意味着只要反向图存活，所有 SavedVariable 引用的中间值都不能被释放。

举个例子：一个几十层的 transformer，每层的反向都要保存若干中间张量（attention 的 QKV 输出、MLP 中间结果、norm 的输入等），每个张量的大小是 batch × seq × hidden 量级 —— 层数与激活值一相乘，反向图存活期间被钉住的显存就是训练显存账单里最大的一块。这也是为什么训练的显存峰值往往出现在 forward 刚跑完、backward 刚开始的附近：所有 SavedVariable 都已经在那里等着了。

activation_checkpoint 的核心思想就是：在前向时把某些区段的 SavedVariable 释放掉，反向到那个区段时重新前向算一遍取回中间值。用计算换显存，对大模型几乎是必备技术。第 18 章 FSDP 章会展开。

7.6 AccumulateGrad：叶子张量的梯度"沉淀池"

leaf 张量的 grad_accumulator 是一个特殊 Node 类：AccumulateGrad（torch/csrc/autograd/functions/accumulate_grad.h）。它的 apply 实现就是把传入的梯度加到张量的 grad_ 字段：

// 简化的 AccumulateGrad::apply
variable_list AccumulateGrad::apply(variable_list&& grads) {
    auto new_grad = grads[0];
    auto* meta = impl::get_autograd_meta(variable);
    if (!meta->grad_.defined()) {
        meta->grad_ = new_grad;        // 第一次设置
    } else {
        meta->grad_ += new_grad;        // 累积
    }
    return {};   // 没有输出 (sink)
}

注意 AccumulateGrad（torch/csrc/autograd/functions/accumulate_grad.h:40，struct TORCH_API AccumulateGrad : public Node）的 next_edges_ 是空的（反向遍历到这里就停了），apply 也不返回任何输出。

torch/csrc/autograd/function.h:89-90 的注释直接把它叫 sink——原文折成两行，先是 "it takes one input, but produces no outputs, instead accumulating"，下一行接 "the input as a side effect"。

7.6.1 为什么训练前要 optimizer.zero_grad()

理解了 AccumulateGrad，你就明白为什么 PyTorch 训练要在每个 batch 开始时调 zero_grad()：

for batch in dataloader:
    optimizer.zero_grad()       # ← 关键
    loss = model(batch).sum()
    loss.backward()
    optimizer.step()

不调 zero_grad 的话，AccumulateGrad 会把当前 batch 的梯度加到上一个 batch 的梯度上，而不是替换。这是 PyTorch 选择"梯度累积式"语义的设计 —— 让"梯度累积"（gradient accumulation，把多个小 batch 的梯度合并成大 batch）这种用法天然支持。

7.6.2 leaf 张量的 grad_accumulator 是 lazy 创建的

回到 §7.2：leaf 张量的 grad_accumulator_ 字段是 weak_ptr<Node>。它什么时候才被实际创建？

答案是第一次有算子需要它的时候。当 c = a + b 触发 VariableType::add，需要为 a 创建 next_edge。代码大致：

// 第一次需要时 lazy 创建
auto a_acc = impl::try_get_grad_accumulator(a);
if (!a_acc) {
    a_acc = std::make_shared<AccumulateGrad>(a);
    impl::set_grad_accumulator(a, a_acc);
}
edge_for_a = Edge(a_acc, 0);

真实实现是 torch/csrc/autograd/variable.cpp 的 impl::grad_accumulator()：加锁后先 grad_accumulator_.lock()，拿到就直接返回，拿不到才 make_shared<AccumulateGrad> 并写回 weak_ptr。所以第二次再用 a 做某算子时，只要上次创建的 accumulator 还活着就会被复用。

这种 lazy 创建让"创建了 leaf 张量但从没用过"的场景不付 AccumulateGrad 的开销 —— 比如 inference 模式下定义了 weight 张量但根本不反向，AccumulateGrad 永远不创建。

7.6.3 AccumulateGrad 的"叶子梯度仅在 leaf 累积"语义

新手常困惑：为什么 c.grad 是 None？c 是 interior 张量，不应该也有梯度吗？

答案是：默认情况下 PyTorch 只在 leaf 张量上累积梯度。这是有意为之的优化 —— 训练里你只关心模型参数（leaf）的梯度，中间张量的梯度只是"传递的中转值"，存下来浪费显存。

如果你确实想看 c 的梯度（如调试），调 c.retain_grad()：

c = a + b
c.retain_grad()              # 告诉 PyTorch 反向时也保留 c.grad
loss = (c * 2).sum()
loss.backward()
print(c.grad)                # 现在有值了

retain_grad 的实现是在 c 身上挂一个 hook，反向流到 c 时把梯度同时存到 c.grad 字段。代价是显存上升（多保留一个张量）。生产代码不该用，调试时很方便。

7.6.4 多次 backward 与 retain_graph

默认情况下 loss.backward() 跑完后反向图被释放（每个 SavedVariable 调 reset_data，shared_ptr 解引用）。如果你想再跑一次反向（如 second-order derivatives），要传 retain_graph=True：

loss.backward(retain_graph=True)   # 保留反向图
# 现在还能再跑
torch.autograd.grad(loss, model.parameters())   # 用 functional 接口再算一次

retain_graph=True 的代价是反向图保留，所有 SavedVariable 仍然持有中间值 —— 这些激活不会随反向推进逐个释放，要一直占到图被丢弃为止。所以只在必要时用，且尽快释放。

第 8 章 §8.5 与 §8.7.1 会展开 GraphTask 怎么处理 retain_graph 标志（keep_graph_）。

7.6.5 optimizer.zero_grad(set_to_none=True) 与 set_to_none=False

老版本 PyTorch 的 zero_grad() 默认行为是 grad_.zero_() —— 一次 inplace 写零。但这要遍历每个参数张量、各调一次 zero kernel，开销不小。set_to_none 参数在 v1.7 加入，默认值在 v2.0 翻转成 True；v2.11 的签名就是 def zero_grad(self, set_to_none: bool = True)（torch/optim/optimizer.py:1028，torch/nn/modules/module.py:2949 同签名），直接把 grad_ 置成 None，省掉那一轮 zero kernel。

set_to_none=True 改变了一个细微语义：之前的"梯度永远是 0 张量"变成"第一次累积前梯度是 None"。这影响一些手工读 param.grad 的代码 —— 它们要先判断 param.grad is None。绝大多数 PyTorch 代码不受影响，但写 grad clip 或自定义 optimizer 时要留意。

省下的量级取决于参数张量的个数与大小（参数越碎、张量越多，省得越明显），需要在自己的模型上实测。

累积式语义的代价就是要手动 zero_grad。新手忘了 zero_grad 会发现损失曲线非常奇怪 —— 这是 PyTorch 训练里最经典的 bug 之一。

7.7 requires_grad / no_grad / inference_mode 的精确语义

第 5 章 §5.6.3 我们对比过这三种 mode，autograd 视角下它们的差异更清楚：

特性	普通张量	no_grad 上下文	inference_mode 上下文
autograd_meta_	nullptr 或 AutogradMeta	同左	通常 nullptr（模式内显式 requires_grad=True 仍会分配）
requires_grad	用户控制	已有张量的标志不变，新算子输出为 False	模式外不能对 inference 张量设为 True
dispatcher 路径	走 AutogradXxx kernel	仍走 AutogradXxx kernel	完全不带 Autograd key
创建反向图	是	否（kernel 内 compute_requires_grad 返回 false 而短路）	否
dispatcher 开销	高 (含 redispatch)	中（少了建 Node，多出来的那一跳还在）	低 (无 autograd key)

inference_mode 少走 autograd key、也不维护 version counter 与 view 元信息，在小算子密集的路径上比 no_grad 更省（算子越小、单次框架开销占比越高，差距越明显；具体数字需自己 profile）。代价是 inference 张量"不可逆"地不能再参与反向 —— 连 SavedVariable 的构造函数里都有一条硬检查："Inference tensors cannot be saved for backward."

实战建议：生产推理服务用 inference_mode（性能优势真实），调试 / 评估用 no_grad（更宽容、可以临时再开 grad）。这种"两档 mode"是 PyTorch 给不同场景留的灵活性。

7.7.1 三种 mode 的实际选择建议

实战里遇到三种 mode 的选择困难，按场景给一份决策表：

场景	推荐
训练循环里只是 forward 看 loss 值	no_grad
训练循环里临时跑一段不参与反向的子图	no_grad
调试 / 单元测试	no_grad
生产 inference server 的 hot loop	inference_mode
Dataloader / 预处理 pipeline	no_grad（产出的张量要喂进训练，第一层就会把它存下来算权重梯度，inference 张量会报错，见下）
只读访问参数（如导出权重、统计参数大小）	inference_mode
在 inference 里偶尔需要反向（如对抗样本生成）	no_grad 配合临时 enable_grad

一条提醒：用 inference_mode 造出来的张量带着 inference 标记，一旦它被喂进需要反向的计算，SavedVariable 会直接报 "Inference tensors cannot be saved for backward"。所以数据 pipeline 用 inference_mode 之前，先确认这批张量不会成为某个反向公式要保存的输入 —— 否则老老实实用 no_grad。

7.7.2 一个不常被提到的细节：grad_mode 与 dispatcher 的协作

with torch.no_grad(): 的实现只是修改 c10::AutogradState 的 thread-local flag（c10/core/GradMode.cpp:9），并不动 dispatch key：keyset 里的 AutogradXxx 照样在，VariableType::mm 这一层照样被 lookup 到并执行，只是它开头的 compute_requires_grad() 撞上 if (!GradMode::is_enabled()) return false;（torch/csrc/autograd/functions/utils.h:60）后走 §7.4.3 那条短路，不建 Node 直接 redispatch。它也不改张量元数据 —— 用户在 with 块内访问张量，仍能看到 requires_grad=True。这种"不改张量、只翻一个线程局部开关"的设计让 grad_mode 成为线程级开关，不影响张量的全局状态。

inference_mode 的实现完全不同：除了在线程局部状态里排除 Autograd 与 ADInplaceOrView 这两层（第 5 章 §5.6.3），它还改张量构造时的 key_set_（c10/core/TensorImpl.cpp 构造函数里的 if (inference_mode) 分支），让构造出来的张量不带 Autograd key。这个区别对调试很重要 —— no_grad 退出后原本的 autograd 行为恢复；inference_mode 创建的张量永久"无法反向"。理解这条区别能避免一些反直觉的 bug。

7.7.3 requires_grad 在 dispatcher 上的体现

requires_grad 并不体现在 key_set_ 上：非 inference 模式下，TensorImpl 构造函数按后端无条件加入 AutogradXxx 等 autograd 相关位（c10/core/TensorImpl.cpp 里还留着 "Ideally we only add AutogradBackend key when the tensor requires grad" 的 TODO），requires_grad 只影响 AutogradMeta。传染性来自 §7.4：compute_requires_grad 发现任一输入 requires_grad，就建 Node 并 set_history，输出有了 grad_fn_，requires_grad() 自然返回 true —— 沿计算图自动传播，用户不需要每次手动设。

7.8 autograd.Function：用户级自定义 autograd

如果你想给 PyTorch 加一个新的可微操作 —— 比如自定义 attention，PyTorch 提供 torch.autograd.Function（Python 接口）和 torch::autograd::Function（C++ 接口）：

class MyExp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.exp()

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * x.exp()

# 使用
x = torch.randn(3, requires_grad=True)
y = MyExp.apply(x)
y.sum().backward()

这套接口的底层实现就是把用户写的 forward 和 backward 包装成一个 Node 类，注入到 PyTorch 的反向图体系。ctx.save_for_backward 先把张量记到 ctx.to_save，forward 返回后再由 C++ 侧（torch/csrc/autograd/python_function.cpp 的 _save_variables）逐个构造 SavedVariable；return grad_output * x.exp() 就是 Node::apply 的返回值。

任何用 autograd.Function 写的算子和 PyTorch 内置算子在反向图层面是平等公民。这套接口让用户能在不改 PyTorch 主仓的前提下扩展 autograd。第 22 章自定义算子会再讲。

7.8.1 autograd.Function 与 torch.library.custom_op 的区别

PyTorch 现在有两套自定义算子接口：

接口	适用场景
torch.autograd.Function	纯 Python 实现，需要自定义反向规则
torch.library.custom_op (v2.4+)	与 dispatcher / torch.compile 完全对接的"主流"接口

autograd.Function 是历史悠久的接口，简单直接。Dynamo 能直接 trace 简单的 autograd.Function（上面的 MyExp 在 v2.11 上用 torch._dynamo.explain 实测 0 次 graph break），但 forward / backward 里一旦有 Dynamo 处理不了的代码，仍会 graph break、回到 eager 路径。

custom_op 是新接口，要求用户多注册几样东西（FakeTensor 实现、autograd 实现等），换来的是 torch.compile 把它当成一个不透明算子处理，行为更可预期。所以新代码推荐 custom_op，旧代码继续用 autograd.Function 也无妨。

第 22 章自定义算子会完整对比这两套接口的优劣与适用场景。

7.8.2 ctx.save_for_backward 的版本检查

ctx.save_for_backward(x) 保存的张量，forward 返回后会被构造成 SavedVariable(x, ...)，记下的是那一刻的版本号。所以 §7.5.1 讨论的版本检查在用户级 autograd.Function 里同样生效 —— 如果你保存了 x，然后在 forward 完成后又对 x 做 inplace 修改，反向时同样会触发"one of the variables needed for gradient computation has been modified"。

反过来，forward 内部在 save_for_backward 之后再 x.add_(1) 并不会触发这道检查（实跑无报错）—— 版本号是 forward 返回后才记的，反向拿到的就是被改过的 x，梯度可能悄悄算错。新手遇到这个报错，常见原因是 Function 的输入或输出在 forward 之后又被 inplace 改了；把 inplace 改成 out-of-place 通常就能解决。确实要在 forward 里原地改输入，要用 ctx.mark_dirty 声明。

7.9 几条容易被忽略的细节

实战里几个易混淆点：

1. tensor.detach() 与 tensor.data 的区别：两者都返回共享 storage、requires_grad=False、grad_fn=None 的新张量，真正的差别在 version_counter，而且方向和很多人的直觉相反 —— detach() 共享原张量的 version_counter（torch/csrc/autograd/VariableTypeManual.cpp 里走 make_variable_non_differentiable_view(..., is_fresh_tensor=true)，内部 data_impl->set_version_counter(impl::version_counter(base))），所以你对 detach 出来的张量做 inplace，原张量的反向会如实报错；而 tensor.data 走的是 VariableHooks::variable_data，它用 shallow_copy_and_detach(/*version_counter=*/0, ...) 换了一个全新的 version_counter 并把 autograd_meta 清空 —— 版本检查就此失效，改了也没人告诉你。这就是".data 是危险接口"的准确含义：不是它绕过 autograd，而是它绕过了那道 inplace 检查

2. torch.set_grad_enabled(False)：等价于 no_grad() 但函数式风格，可以在条件分支里更灵活地控制

3. tensor.requires_grad_() 只能改 leaf：对 leaf 张量随意调；对 interior 张量（有 grad_fn）调 requires_grad_(True) 是空操作，调 requires_grad_(False) 会报 "you can only change requires_grad flags of leaf variables."，因为这会破坏反向图一致性（报错里建议改用 detach()）

4. with torch.no_grad(): 内部创建的张量：这些张量 grad_fn 全是 None。即便 with 块外面再设 requires_grad=True 也无济于事 —— 反向图被永久砍掉了

5. tensor.retain_grad()：默认情况下 interior 张量没有 grad_ 字段（反向时梯度只流向 leaf）。如果你想看某个中间张量的梯度，要在前向时调 retain_grad()。它的实现是给那个 interior 张量挂一个特殊 hook，反向流到这一步时把梯度同时存到 grad_ 字段里。代价是显存上升

6. autograd 与多 device 张量：前向里的跨设备搬运（.to() / .cuda() / .cpu() 走 _to_copy）会留下一个 ToCopyBackward0 节点，反向时由它把梯度拷回原设备（tools/autograd/derivatives.yaml 里 _to_copy 的公式是 _to_copy_backward）。引擎本身不做隐式搬运：梯度所在设备与节点记录的不符时，只有 0 维张量会被自动 .to()，其余直接报 "invalid gradient at index ..."（torch/csrc/autograd/engine.cpp 的 validate_outputs_impl）。所以反向里的每次 H2D / D2H 都对应前向里一次搬运，这很方便，但容易隐藏性能问题 —— 用 profiler 找到这些拷贝是优化训练吞吐的常见动作

7. with torch.enable_grad(): 嵌套在 no_grad() 里：这是合法的，让你在 inference 大块代码里临时开启 autograd，比如在 no_grad 的评估代码里临时对输入求梯度（对抗样本、显著性图）

7.10 backward 时的钩子 (hook) 系统

autograd 还支持在反向流程的特定时刻插入用户回调：

Hook 类型	注册方式	触发时机
Tensor backward hook	tensor.register_hook(fn)	反向流到这个张量时（grad 已经算出但还没传给 next_edges）
Module hook	module.register_full_backward_hook(fn)	反向经过整个 nn.Module 时
Post-accumulate-grad hook	param.register_post_accumulate_grad_hook(fn)	leaf 张量的 grad 累积完成之后（AutogradMeta::post_acc_grad_hooks_）

最常见的用法是梯度裁剪 (gradient clipping)：

# 经典用法：裁剪超过阈值的梯度
for p in model.parameters():
    if p.requires_grad:
        p.register_hook(lambda grad: torch.clamp(grad, -1.0, 1.0))

但生产代码通常不这样用 —— 因为 torch.nn.utils.clip_grad_norm_ 提供了更标准的接口。Hook 主要在调试时有用：在某个层的反向时打印梯度统计、检查 NaN 是否在某层出现等。

这套 hook 机制的实现是在 Node 上挂四组容器（torch/csrc/autograd/function.h:699-703）：pre_hooks_、tensor_pre_hooks_（两者都是 vector<unique_ptr<FunctionPreHook>>）、按输出下标索引的 retains_grad_hooks_，以及 post_hooks_（vector<unique_ptr<FunctionPostHook>>）。反向流到那个 Node 时按类别与顺序依次调用。接口定义在 torch/csrc/autograd/function_hook.h。

7.11 跨专栏关联

• 《Tokio 源码深度解析》专栏第 5 章（多线程 Scheduler 与工作窃取）：autograd Engine 同样是"任务队列 + 常驻 worker 线程"，但它按 device 路由任务、不做 work-stealing，和 Tokio 的工作窃取调度器既相似又不同。第 8 章 §8.11 会做详细对照
• 《MCP 协议设计与实现》专栏第 3 章（JSON-RPC 与消息格式）：MCP 在分布式调用里传递 context 与 PyTorch 在 dispatcher 里通过 c10::DispatchKeySet + thread-local 传递 autograd 状态有相通思想

7.12 一个练习：手画一段代码的反向图

import torch

a = torch.randn(3, requires_grad=True)
b = torch.randn(3, requires_grad=True)
c = a * b
d = c.sum()
e = a.exp()
f = e.sum()
loss = d + f

试着画出反向图：每个张量是哪种类型（leaf / interior）、每个 grad_fn 是什么 Node 类、每条 next_edge 指向哪里、a 的梯度怎么从 loss 流回来。

提示：a 出现在两条路径上（a*b 和 a.exp()），所以 AccumulateGrad_a 会有两条 in-edge。loss.backward() 时这两路梯度会自动求和到 a.grad。

画完之后用 torch.autograd.grad 或者 loss.backward() 验证：

loss.backward()
print(a.grad)
# 应该是 b + a.exp()  (因为 d=a*b, f=a.exp().sum, loss=d+f → dL/da = b + e^a)

7.13 一些 autograd 关键细节的总结

把 autograd 的精髓压成几条：

1. autograd 是 dispatcher 的中间层：通过 Autograd* DispatchKey 注入到正常算子调用链，redispatch 后让真正的 backend kernel 接手
2. 反向图在前向时建好：每个可微算子调用都创建一个 Node，next_edges 指向上游 Node，让整张图自然链起来
3. Node 是 base class，每个算子有自己的 XxxBackward 子类：apply 方法实现具体反向数学
4. leaf 张量有 grad_accumulator (AccumulateGrad)，interior 张量有 grad_fn，两者互斥
5. SavedVariable 用 version_counter 防 inplace 污染，保存输出时存不带 grad_fn 的浅拷贝（inplace-on-view 时用 weak_ptr）防循环引用
6. requires_grad 沿计算图自传染，no_grad / inference_mode 提供两档"关闭 autograd"的语义

如果你能在脑子里画清"a + b 触发的 dispatcher → VariableType::add → 创建 AddBackward0 → next_edges 指向 a/b 的 grad_fn → redispatch 调真 add → set_history → 返回 c"这一连串过程，本章就内化了。

7.14 Compiled Autograd：把反向图也编译

Compiled Autograd（torch/_dynamo/compiled_autograd.py，由 PR #103822 引入，首次随 v2.1.0 发布，此后一直在演进）让反向 Engine 也能被 Inductor 编译。简单原理：

1. 用户照常 forward + loss.backward()
2. PyTorch 把"按 next_edges 反向遍历 + 调每个 Node.apply"这套调度逻辑自身捕捉成 FX graph
3. 把这个 graph 交给 torch.compile 的后端（默认 Inductor，GPU 上生成 Triton kernel）编译
4. 后续 backward 直接调编译后的代码，跳过 Engine 逐 Node 调度的开销

收益主要来自两处：一是把 Engine 的逐 Node 调度换成一段编译好的代码，二是让 Inductor 有机会跨 Node 融合反向 kernel —— 所以算子越小、Node 越多，收益越明显（LSTM、MoE 这类反向是典型受益者）。代价是首次要付编译时间，之后命中缓存。具体倍数没有普适答案，要在自己的模型上量。

启用开关就在 dynamo 的配置里（torch/_dynamo/config.py:719 定义 compiled_autograd = False）：

# 调用侧用法示例，不是 PyTorch 源码
import torch._dynamo
torch._dynamo.config.compiled_autograd = True

第 13 章 AOTAutograd 章会展开。理解 compiled autograd 的工作原理需要先吃透本章（autograd 怎么建反向图）+ 第 8 章（engine 怎么执行）+ 第 12 章（Dynamo 怎么 trace）—— 它是这几个系统协同的产物。

7.15 几条工程经验

实战里 autograd 相关 issue 大致分三类，附常见诊断方法：

类型 A：反向报错 "one of the variables ... has been modified by an inplace operation"

诊断：用 torch.autograd.set_detect_anomaly(True) 让 PyTorch 在前向时记录每个 Node 的创建栈，反向报错时打印精确位置。修法：把可疑的 inplace 操作（*=、+=、add_、relu_）改成 out-of-place。

类型 B：训练 loss 不下降 / NaN 不收敛

诊断：可能是忘了 zero_grad 让梯度累积。或者某个 Node 的反向规则有 bug 导致梯度方向错。用 torch.autograd.gradcheck 可以数值验证某个算子的反向是否正确。

类型 C：显存峰值过高

诊断：通常是 SavedVariable 持有太多中间值。用 torch.cuda.memory._record_memory_history() snapshot 看（第 4 章 §4.11）。修法：activation_checkpoint、混合精度（第 20 章）、CPU offload（saved_tensors_hooks）。

7.16 跨专栏关联补充

• 《vLLM 推理内核深度解析》专栏第 4 章（PagedAttention）：vLLM 的推理路径完全不构建反向图（pure inference）。理解 PyTorch autograd 的"非侵入式"设计后再看 vLLM 跳过 autograd 的方式，就能理解 PyTorch 在训练 / 推理两侧的统一性
• 《Serde 元编程》专栏第 12-13 章（派生宏代码生成）：tools/autograd/gen_variable_type.py 自动从 derivatives.yaml 生成 VariableType::add 等包装函数，与 Serde 派生宏（自动生成 serialize/deserialize 代码）思路相通

7.17 几条 autograd 设计的"通用启示"

如果你设计自己的 AD 系统（深度学习外的领域：金融衍生品定价、物理模拟），本章思想能直接迁移：

第一：反向图作为 DAG，每个节点写自己的 apply —— 不要试图把整张图集中在一个数据结构里管理。每个 Node 自治、用 Edge 连接，是最 modular 的设计

第二：分离声明（forward）与执行（backward） —— forward 时只建 IR，不真做反向计算。这让 forward 的代码路径专注于生成图，backward 路径专注于调度执行

第三：用 weak_ptr 打破循环引用 —— 任何"图节点 + 张量元数据互引"的设计都要小心循环引用，weak_ptr 是经典武器

第四：version_counter 防 inplace 污染 —— 任何"延迟计算"系统都要有"我保存的值是不是还有效"的校验机制，version 号是最简方案

第五：lazy materialize_autograd_meta —— 元数据按需分配。绝大多数对象不需要那块元数据，给所有对象都分配是巨大浪费

第六：version_counter / weak_ptr / lazy 创建是性能预算的"三件套" —— 任何"延迟构造 + 多对象共享 + 安全检查"的系统都会用到这三件套，PyTorch 是它们配合极致的范例

把这六条记下来，写自己的 AD 系统能少走很多弯路。

7.18 一个反思：为什么 PyTorch 选了"运行时建图"

最后回到一个高维度问题：为什么 PyTorch 选择前向时偷偷建反向图这条路，而不是 JAX 那种"先 trace 再编译"或者 TF 1.x 那种"显式定义图"？

PyTorch 设计者的回答（参考 NeurIPS 2019 论文 PyTorch: An Imperative Style, High-Performance Deep Learning Library）：

• define-by-run 让控制流自由：用 Python 的 if/for 写动态网络，反向图自然跟着控制流走，不需要 tf.cond / lax.scan 这类图原语
• 构图开销可承受：每个算子建一次 Node + 几条 Edge 只是少量对象分配与指针连接，相对于大张量上的数学计算通常可以忽略
• 代码即调试器：反向图等同于"前向跑过的代码路径"，调试时打开 stack trace 立刻看到结构

代价是反向图每次都重建（不能像静态图那样跨 step 复用编译）。torch.compile 在 v2.0 之后用"AOTAutograd"机制把反向图也编译进 Inductor，弥补了这个性能缺口（第 13 章会展开）。这种"动态图 + 可选编译"是 PyTorch 哲学的核心，也是它赢得学术界 → 工业界飞轮的根本。

值得对比的是：JAX 的 grad 是函数变换 g = jax.grad(f)，把整个 f 变换成一个新的梯度函数（再套 jax.jit 才编译）。它要求 f 是纯函数（不能有 inplace、不能有副作用），换来的是整段反向函数都能交给编译器优化。PyTorch 的 loss.backward() 是引擎驱动的图遍历，允许 inplace 与各种动态行为，代价是反向时不能像 JAX 那样跨调用复用编译。

两条路在哲学上是镜像的：JAX 把"什么是反向"拆成数学问题让编译器解；PyTorch 把"反向怎么调度"拆成数据结构让引擎跑。前者依赖严格函数式假设，后者依赖运行时元信息。理解这两种 AD 设计的本质差异，你能在选型时做出有依据的判断 —— 学术研究 / 动态网络偏 PyTorch，函数式建模 / 静态优化偏 JAX。

下一章拆 Engine：这张反向图怎么被多线程引擎执行 —— 看 PyTorch 按 device 路由的 ReadyQueue 调度。