PyTorch 训练框架内核深度解析

第13章 AOTAutograd：函数化与正反向图划分

作者 杨艺韬 · 13,588 字 · 发布于 · 更新于

Functions in this module do most of the "work" of AOTAutograd.

—— torch/_functorch/_aot_autograd/graph_compile.py 模块 docstring 第一句。紧接着它列出 aot_dispatch_* 函数的五个步骤：接收 flat_fn、flat_args 和一批 metadata，跑一组编译前包装（比如参数去重），跑真正的编译器，把返回的 callable 再套上一组编译后包装，最后连同 metadata 一起返回

本章要点

  • AOTAutograd 在 Dynamo 之后、Inductor 之前：把 forward-only graph 扩展成 (forward + backward) joint graph
  • Functionalization 把 inplace 操作改写成纯函数式：x.add_(y) → x = x.add(y)，让下游 IR 没有 mutation 副作用
  • min-cut partition 用图论里的最小割算法决定"哪些中间张量保留给反向用"：在显存与重算之间找最优平衡
  • 三个产物：forward graph + backward graph + saved tensor list
  • AOTAutograd 不绑定 Inductor：aot_eager backend 只跑 AOTAutograd、不上 Inductor，等同于"trace 反向但不编译"，用于调试
  • 它是 PyTorch 编译器栈最复杂的一段：torch/_functorch/ 整个 namespace 三万余行 Python，光 _aot_autograd/ 就有一万六千多行，处理无数 corner case（subclass、view、inplace、AMP、checkpoint）

13.1 Dynamo 留下的烂摊子

第 12 章我们看到 Dynamo 输出 forward FX Graph + Guards + example_inputs。但 Inductor 想要的更多：

  • Joint forward-backward graph：训练时 backward 也是热路径，必须一起编译
  • 纯函数式 IR：Inductor 内部图变换不希望处理 inplace 副作用
  • 明确的 saved tensor 列表：哪些中间值要保留给 backward

这三件事都不在 Dynamo 的职责内。AOTAutograd（Ahead Of Time Autograd）补上这一段。源码主要在 torch/_functorch/aot_autograd.py（v2.11 实测 1878 行）和 torch/_functorch/_aot_autograd/ 目录（实测 ~16500 行）。整个 torch/_functorch/ namespace 共 ~32000 行，是 PyTorch v2.x 编译栈的核心组件。

13.2 整体流程

flowchart TB
    Dy[Dynamo 输出<br/>forward FX Graph]
    Dy --> Trace[1 反向 trace<br/>用 autograd.grad 在 fake tensor 上跑]
    Trace --> Joint[joint graph<br/>forward + backward 节点混在一起]
    Joint --> Func[2 functionalization<br/>inplace → 纯函数]
    Func --> Part[3 min-cut partition<br/>切成 fw_module + bw_module]
    Part --> FW[forward 子图]
    Part --> BW[backward 子图]
    Part --> Save[saved tensors 列表]
    FW --> Ind[Inductor 编译 forward]
    BW --> Ind2[Inductor 编译 backward]
    style Joint fill:#fef3c7,stroke:#f59e0b
    style Part fill:#dbeafe,stroke:#3b82f6,stroke-width:2px

四步：trace → functionalize → partition → 各自送 Inductor。（源码里前两步其实是同一趟 make_fx：torch/_functorch/_aot_autograd/graph_capture.py 先用 create_functionalized_fn 把 joint 函数包好再交给 make_fx，边 trace 边函数化；分开画是为了讲清楚。）

13.3 反向 trace：用 fake tensor 跑出 backward

Dynamo 给的是 forward graph（一组 torch 层的函数 / 方法调用，如 l_x_ @ l_w_、matmul.relu()，还不是 ATen 算子）。AOTAutograd 怎么得到 backward？答案是 再 trace 一遍：

def joint_fn(primals, tangents):
    out = forward_graph(*primals)
    grads = torch.autograd.grad(out, primals, tangents)
    return out, grads

primals 是输入张量，tangents 是输出梯度。autograd.grad 触发反向 —— 但所有张量都是 FakeTensor（第 5 章 §5.7 提过 FakeTensor 是只算 shape/dtype 的"假张量"）。整个反向不真做计算，只生成节点。

trace 完拿到一张 joint graph：包含 forward 节点 + backward 节点 + 所有中间值。这张图就是 后续两步的输入。

13.4 Functionalization：消除 inplace

joint graph 里可能有 x.add_(y)、x.copy_(y) 等 inplace 操作。这种操作让图分析变难（节点 X 的输出可能被节点 Y 偷偷修改）。Functionalization 把它们重写成纯函数式：

inplace 版:  x.add_(y)             # x 被原地修改
functional: x_new = x.add(y)        # x 不变, 返回新张量

实现机制：用 FunctionalTensorMode（第 5 章 §5.7 讲的 TorchDispatchMode 之一）拦截每个 inplace 算子，把它替换成对应的 out-of-place 版本 + 一个 copy_back（如果用户原本期待原地修改）。这套改写对用户透明，但下游 IR 里所有节点都变成纯函数。

为什么这件事重要？因为 Inductor 的算子融合 / 重排算法假设节点没有副作用。有 inplace 时，"a + b 节点"的输入可能在它执行前被另一个节点 inplace 改了 —— 重排顺序就崩。Functionalization 一刀斩断 mutation，让下游图变换可以放心做。

代价是显存上升（每个 inplace 变成 out-of-place 多一份张量）。但 Inductor 后面有 buffer reuse pass 能把这些"逻辑上独立但实际可复用"的张量重新合并 inplace —— 等同于先抹掉 inplace 让分析简单、再恢复 inplace 让运行高效。这种"先简化再优化"是编译器经典做法。

13.5 min-cut partition：在显存与重算之间找最优

joint graph 里 forward 与 backward 节点混在一起。要切成两段子图，关键问题是：哪些 forward 中间张量要保留给 backward 用？

直观选项有两个极端：

  • 全保留：所有 forward 中间值都保留 → 显存爆炸
  • 全重算：backward 时从 input 重新跑一遍 forward → 计算翻倍

实际最优在两者之间：保留"算起来贵的"中间值，丢弃"重算便宜的"。这就是经典的最小割问题。

torch/_functorch/partitioners.py:min_cut_rematerialization_partition（v2.11 中位于该文件 line 3449 的核心函数）建模为：

  • 每个 joint 节点拆成 X_in / X_out 两个点，中间那条边的容量是该节点输出的字节数（保留代价 = 显存），再按离反向的距离、是否会被物化做加权
  • "重算贵"不进边权，而是直接禁掉：matmul、随机算子等不许重算的节点（§13.9）从 source 连一条无穷容量边；反向必需的节点连到 sink
  • 用网络流算法（max-flow / min-cut）找出"切割集"，让被割节点的总字节数最小

切割集就是要保留的张量集合。算法保证 forward 之后只把这些张量传给 backward，其他全部丢弃；backward 时从这些保留的张量出发，必要时重算。

这套算法的优雅之处：它自动决定哪些中间值存、哪些重算，用户不需要手动标"这层重算"。但默认它只重算白名单里的便宜算子（pointwise、类型转换、view 等），matmul 这类算子的输出一律存下，所以它不等于把整层 checkpoint 掉。

在 transformer 上，min-cut 的典型取向是保留"重算贵、输出小"的中间值，丢弃"重算便宜"的——这与人手设计 activation_checkpoint 时的直觉一致。具体某个模型上留了什么、丢了什么，可以开 TORCH_LOGS=aot_joint_graph,aot_graphs 对照切分前的 joint graph 和切分后的 fw / bw 图（partitioner 不是合法的 TORCH_LOGS 名字，写了会直接报错）。

13.6 输出三件套

partition 完拿到：

  1. fw_module：forward 子图 —— 接收 primals，返回 outputs + saved tensors
  2. bw_module：backward 子图 —— 接收 saved tensors + tangents，返回 gradients
  3. saved tensors 的位置信息 —— 告诉 runtime 哪些 fw 输出要喂回 bw。源码里不是一张索引表，而是 ViewAndMutationMeta 算出的下标与切片（num_forward_returns 在 torch/_functorch/_aot_autograd/schemas.py:644，tensors_saved_for_backwards_slice 在 :707）

Runtime 流程示意（forward 调用时；saved_tensor_indices / real_output_indices 是示意用的名字）：

fw_outputs = fw_module(*primals)
saved = [fw_outputs[i] for i in saved_tensor_indices]
real_outputs = [fw_outputs[i] for i in real_output_indices]
# 用户拿到 real_outputs
# saved 暂存等 backward 用

# backward 调用时
gradients = bw_module(*saved, *tangents)

这套 runtime 由 torch/_functorch/_aot_autograd/runtime_wrappers.py（3034 行）实现。它包了一层 autograd.Function，让编译产物对接到 PyTorch 的 backward Engine（第 8 章）—— 从用户角度看跟普通 module 一样，loss.backward() 自动触发 bw_module。

13.7 runtime_wrappers 的层叠

torch/_functorch/_aot_autograd/runtime_wrappers.py（3034 行）有十几个 Wrapper 类，每个负责一类问题。它们像洋葱一样层层包裹编译产物：

Wrapper 行号 职责
RuntimeWrapper 165 包在 autograd.Function 外面的运行时收尾：把留在图外的 input mutation 写回输入、重新生成与输入别名的输出、处理 detach 与 AMP 状态
FunctionalizedRngRuntimeWrapper 731 处理 RNG state（让随机算子也能 functionalize）
FakifiedOutWrapper 813 按编译时记录的 stride 修正输出张量的 stride
AOTDispatchSubclassWrapper 906 tensor subclass（DTensor、TwoTensor 等可 trace 的 wrapper subclass）的 unwrap / rewrap
EffectTokensWrapper 993 side-effect 算子（如 print、torchbind 调用）的 token 串联
AOTDedupeWrapper 1106 去重：同一张量作为多个 input 时合并
AOTSyntheticBaseWrapper 1356 view 合成 base：多个 view 共享 storage 时统一处理
AOTDispatchAutograd 2303 把 fw + bw 包成 autograd.Function 接入 Engine
DebugAssertWrapper 2940 debug 模式下的运行时断言

每一层解决一个具体 corner case。运行时调用 compiled_fn(*args) 时，从外到内穿过所有 wrapper。下图是训练路径的套法，按 graph_compile.py 里的组装顺序画（_aot_stage2c_make_autograd_function 与 _aot_stage2b_compile_forward_or_inference；DebugAssertWrapper 只在 config.debug_assert 打开时插在 SyntheticBase 与 RuntimeWrapper 之间）：

flowchart LR
    User[用户调用 compiled_fn]
    User --> R1[AOTDedupeWrapper<br/>去重]
    R1 --> R2[AOTSyntheticBaseWrapper<br/>view 合并]
    R2 --> R3[RuntimeWrapper<br/>图外 mutation / 别名收尾]
    R3 --> R4[AOTDispatchAutograd<br/>CompiledFunction.apply]
    R4 --> R5[FakifiedOutWrapper<br/>修正输出 stride]
    R5 --> R6[FunctionalizedRngRuntimeWrapper]
    R6 --> R7[AOTDispatchSubclassWrapper<br/>unwrap subclass]
    R7 --> R8[EffectTokensWrapper<br/>串 side effect]
    R8 --> Real[编译后的 fw_module]
    style R4 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

R5–R8 包的是 forward 编译产物，在 CompiledFunction.forward 里调用；反向时 CompiledFunction.backward 另行调用编译后的 bw_module。

这种"每个 wrapper 一类问题"的设计是 _aot_autograd/ 一万六千多行代码却仍可读的关键。读源码时按层定位 —— 看到具体错误信息搜对应 Wrapper 类名能精准找到出问题的层。

13.8 AOTConfig：编译策略的中央配置

torch/_functorch/_aot_autograd/schemas.py:1066 的 AOTConfig 是 AOTAutograd 的配置 hub：

@dataclass
class AOTConfig:
    fw_compiler: Callable          # forward backend (通常是 Inductor)
    bw_compiler: Callable           # backward backend
    partition_fn: Callable          # torch.compile 路径下是 min_cut_rematerialization_partition；aot_function 默认 default_partition
    decompositions: dict            # ATen op 拆解规则 (大 op 拆成小 op)
    num_params_buffers: int
    aot_id: int
    keep_inference_input_mutations: bool
    dynamic_shapes: bool
    # ... 共二十多个字段（均为节选，类型标注有简化）

几个特别值得注意的字段：

  • decompositions：把"难编译的大 op"拆成"易编译的小 op 组合"。比如 Inductor 的表会把 aten::native_layer_norm 拆成 var_mean + add + rsqrt + sub + mul + add 等（§13.10 有实跑结果）。每个 backend 自己决定传什么表：Inductor 传自家的表，aot_eager 不传（不拆）。decomp 让 AOTAutograd 给 backend 一致的"低级 IR"
  • keep_inference_input_mutations：决定 inference 图里 input mutation 落在哪里。为 True 时 copy_ 直接留在图末尾；为 False 时被改输入的新值作为额外输出返回，由 RuntimeWrapper 在图外写回。两种方式 mutation 都会生效，也都要先 functionalize（§13.13 有实跑对比）
  • dynamic_shapes：shape 是符号还是具体值。影响整个 trace + partition + codegen

AOTConfig 像 OpenSSH 的 ssh_config：所有 knob 集中在一个对象，调编译策略时改这里。

13.9 min-cut 的 banned / recomputable 列表

partitioners.py:must_recompute 与相关函数控制哪些节点必须 recompute、哪些必须 save。具体规则（节选，v2.11 默认配置）：

Banned (必须 save，不能 recompute)：should_ban_recomputation 返回原因串的几类

  • 标记为 CheckpointPolicy.MUST_SAVE 的（用户 activation_checkpoint 显式标记的，"marked MUST_SAVE"）
  • 不在可重算白名单里的算子（ban_recompute_not_in_allowlist 默认 True，"not in recomputable allowlist"）：aten.mm / aten.addmm / 卷积 / SDPA 这些 compute_intensive_ops，以及 aten.native_dropout 等随机算子（recompute 会得到不同值）都不在白名单里
  • 结果在反向里必须物化的（"materialized in backward"）
  • 输出比输入小 4 倍以上的归约（"reduction op"，存它很便宜）
  • 离反向太远的（"too far from backward"）
  • 此外 primal 输入与 _c10d_functional 命名空间的 collective 也不重算

Force recompute (必须 recompute，不能 save)：

  • 用户标记 CheckpointPolicy.MUST_RECOMPUTE 或 PREFER_RECOMPUTE 的（must_recompute 两种都认）
  • 打开 recompute_views=True（默认 False）时的 view 算子：权重设为无穷大，永不保存

Negotiable (min-cut 算法决定)：

  • 剩下的白名单算子——pointwise、prims.convert_element_type / aten._to_copy 这类类型转换、view 等"重算便宜"的节点。min-cut 按输出字节数（再按离反向的距离、是否物化加权）决定存哪几个张量

这套 banned / recomputable / negotiable 三态分类让 min-cut 算法只在合理空间搜索，不会试图重算 matmul 浪费算力、也不会试图 recompute 随机算子破坏正确性。

源码里的实现入口是：torch/_functorch/partitioners.py:93 的 OpTypes 数据类（把算子按 fusible_ops / compute_intensive_ops / random_ops / view_ops / recomputable_ops 分类）、:162 的 must_recompute(node)（读 node.meta["recompute"] 里的 CheckpointPolicy），以及 :2062 的 should_ban_recomputation(node)（返回"为什么这个节点不能重算"的原因串，如 "marked MUST_SAVE"）。

13.10 decomposition table：把大算子拆成小算子

AOTConfig.decompositions 是把 ATen "大算子"自动拆成若干"小算子组合"的规则表。源码主仓在 torch/_decomp/decompositions.py（5490 行），用 @register_decomposition(aten.xxx) 装饰器批量注册。

为什么要拆？因为 Inductor 想要看到的 IR 比 ATen 更"基础"。比如 aten.native_layer_norm 是一个复合算子（含 mean / var / sub / div / mul / add），如果直接送给 Inductor，它得给每个复合算子写一份 lowering 规则，工作量爆炸。decomposition 把 layer_norm 拆成几个简单 op 后，Inductor 只需要 lower 这几个简单 op，整个 layer_norm 自动用上 fusion / autotune 机会。

具体例子（精简版；native_layer_norm 这一条实际注册在 torch/_refs/__init__.py:3410，_refs 里的参考实现也经同一个装饰器进表）：

@register_decomposition(aten.native_layer_norm)
def native_layer_norm(x, normalized_shape, weight, bias, eps):
    var, mean = torch.var_mean(x, dim=..., unbiased=False, keepdim=True)
    rstd = torch.rsqrt(var + eps)
    out = (x - mean) * rstd
    if weight is not None:
        out = out * weight
    if bias is not None:
        out = out + bias
    return out, mean, rstd

——一行 aten.native_layer_norm 调用被拆成一串基础算子。本机用 backend="aot_eager_decomp_partition"（套用 Inductor 的 decomp 表、但不上 Inductor）对 F.layer_norm(x, (8,), w, b) 实跑，forward 图是 var_mean → add → rsqrt → sub → mul → mul → add 这 7 个 ATen 算子。Inductor 的 fusion 算法可以把它们合并进一个 kernel，并能与周边算子继续 fuse（本机无 GPU，按源码判断、未实跑）。

torch/_decomp/__init__.py:291 的 core_aten_decompositions() 返回 PyTorch 维护的"标准 decomp table"：显式列出的约 210 个算子（441 个 overload，_core_aten_decompositions_post_autograd()），再自动补上 CompositeImplicitAutograd 算子的拆法，本机实测共 999 个条目。它是 export 的 run_decompositions() 的默认表；Inductor 的表（torch/_inductor/decomposition.py:110）以它为底再叠加自家规则；直接调 aot_function 不传 decompositions 时则是空表。第三方后端可以只用部分 decomp 或加自家规则：

import torch
from torch._decomp import get_decompositions

aten = torch.ops.aten
# 我家硬件原生支持 layer_norm, 不要拆
my_decomps = get_decompositions([aten.relu, aten.silu])  # 按需列出, 不含 layer_norm

decomp table 这套机制是 PyTorch 高层 API 与编译器后端解耦的关键：用户写 torch.layer_norm，AOTAutograd 拆成 primitives，每个后端选自家友好的拆法。这与第 5 章 §5.2.2 的 CompositeImplicitAutograd 思想一致 —— 但前者是编译期拆、后者是运行期拆。

理解了 decomp table，你就理解了为什么 Inductor 能在不写"layer_norm 专用 codegen"的前提下生成高质量的 layer_norm kernel —— 拆 + fusion 自动覆盖。这条思想是编译器栈"少写专用代码、多用通用算法"的工程范本。

13.11 三种 dispatch 入口：autograd / export / inference

AOTAutograd 不是单一入口，它有三种主要 trace 模式：

v2.11 里这三条路径被组织成"stage1 抓图 + stage2 编译"两段。抓图阶段在 torch/_functorch/_aot_autograd/graph_capture.py：aot_dispatch_base_graph（:160，只要 forward）与 aot_dispatch_autograd_graph（:383，要 joint）。编译阶段在 graph_compile.py：

  • aot_stage2_autograd（:2159）：训练路径。joint graph 走 partition，fw / bw 分别给 Inductor。最常用
  • aot_stage2_inference（:422）：推理路径。只有 forward，跳过 partition
  • aot_stage2_export（:269）：export 路径（aot_export_module 等入口）。不编译，直接把 fx 图交回调用方；第 15 章 §15.10.1 的 AOTI 在 torch/_inductor/compile_fx.py 里就是先调 aot_export_module(..., trace_joint=False) 走这条

三者共享绝大部分代码，差异在最后阶段：

阶段 autograd inference export
trace joint ✓ ✗ (只 forward) ✗ (只 forward)
functionalize ✓ ✓ ✓ (更严格)
min-cut partition ✓ ✗ (无 backward 可分) ✗
接 autograd.Function ✓ ✗ ✗
dynamic shape 支持 支持 需显式声明

理解三种入口让你看 PyTorch 源码时不困惑：相同的抓图流程被三个 stage2 函数复用（aot_stage1_graph_capture 在 torch/_functorch/_aot_autograd/graph_compile.py:173，是三者共同的上游）。autograd 与 inference 之间怎么选，不由 Dynamo 决定，而是 AOTAutograd 自己在 torch/_functorch/aot_autograd.py 的 create_aot_state 里判断：有输入 requires_grad 就先记 needs_autograd，若发现没有输出需要梯度、也没有需要梯度的输入被原地修改，再降回 inference；随后 aot_stage2_compile（torch/_functorch/_aot_autograd/graph_compile.py:343）据此分派。

13.12 FunctionalTensor 与 FunctionalTensorMode

§13.4 提了 functionalization 用 FunctionalTensorMode 拦截 inplace。具体实现：

torch/_subclasses/functional_tensor.py:59 的 FunctionalTensor 是 torch.Tensor 子类，它是 C++ 函数化逻辑外面的一层薄 Python 壳：每个实例的 elem 字段是一个 C++ FunctionalTensorWrapper（aten/src/ATen/FunctionalTensorWrapper.h，内部用 generation_ 计数追踪 base 的更新）。inplace op（如 x.add_(y)）被 FunctionalTensorMode（:335）拦截：

# 简化版拦截逻辑
class FunctionalTensorMode(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args, kwargs):
        if is_inplace(func):
            # 把 inplace 改成 out-of-place
            new_value = ops.add(x, y)              # 不 inplace
            x.replace_(new_value)                   # 换掉 wrapper 持有的底层 (functional_tensor.py:266)
            x.commit_update()                       # 登记这次更新 (:269)
            return x
        ...

replace_ + commit_update 让 FunctionalTensor 看起来像被 inplace 修改，底层张量其实已经换了（trace 期间它没有真实存储，调 x.data_ptr() 会直接抛 RuntimeError，本机实测）。外部代码看到的语义与原始 inplace 一致，但内部已经是纯函数式。

这套机制的工程精妙之处：用户原本写 x.add_(y) 期望"x 被原地改"。functionalize 后底层 ATen 看到的是 x = x.add(y)，但 x 这个 wrapper 仍指向更新后的值 —— 用户层面没感知。

13.13 copy_back 机制：functionalize 后处理 mutation

functionalize 把 inplace 全改 out-of-place，但有些场景用户确实需要 mutation 生效（如训练时 BN 的 running_mean.copy_(new_mean)、optimizer 的 param.data.copy_(...)）。

AOTAutograd 的解法：在 graph 末尾插入 copy_ 节点把 functionalize 后的新值写回原位置：

原始: param.data.add_(grad, alpha=-lr)
functionalize 后:
  new_param = param.add(grad, alpha=-lr)     # out-of-place
  ... (更多计算)
  param.copy_(new_param)                       # 末尾的 copy_back

copy_back 让"用户原本期望的 inplace"在 graph 末尾生效。Inductor 看到这种 copy_ 时不会再 functionalize（已经在 graph 末尾、没下游使用）—— 直接生成内存拷贝指令。

AOTConfig.keep_inference_input_mutations（§13.8）控制这次写回放在图内还是图外，不管哪种，mutation 都会生效。下面是本机对 p.add_(g, alpha=-0.1); return (p * 2).sum() 调 aot_function 的实跑结果：

# keep_inference_input_mutations=False（aot_function 的默认）：新值作为额外输出, 由 RuntimeWrapper 在图外写回
def forward(self, arg0_1, arg1_1):
    add = torch.ops.aten.add.Tensor(arg0_1, arg1_1, alpha = -0.1);  arg0_1 = arg1_1 = None
    mul = torch.ops.aten.mul.Tensor(add, 2)
    sum_1 = torch.ops.aten.sum.default(mul);  mul = None
    return (add, sum_1)

# keep_inference_input_mutations=True（Inductor、aot_eager 传的都是 True）：copy_ 留在图末尾
def forward(self, arg0_1, arg1_1):
    add = torch.ops.aten.add.Tensor(arg0_1, arg1_1, alpha = -0.1);  arg1_1 = None
    mul = torch.ops.aten.mul.Tensor(add, 2)
    sum_1 = torch.ops.aten.sum.default(mul);  mul = None
    copy_ = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = add = copy_ = None
    return (sum_1,)

两次调用后 p 都等于 p - 0.1 * g。即便开了这个 flag，input mutation 能否留在图内还另有条件（例如被改的输入 requires_grad、且修改不在 no_grad 下，就只能放图外），判断在 torch/_functorch/_aot_autograd/functional_utils.py:541 的 _check_if_mutation_can_be_in_graph。

13.14 min-cut 算法的网络流实现

§13.5 讲了 min-cut 的"图论建模"。具体怎么解？partitioners.py 用 NetworkX 库的 max_flow 算法。

建图：

  • 源 source 节点：连向 primal 输入和所有被 ban 的节点（无穷容量，保证它们不会被重算）
  • 汇 sink 节点：由反向必需的节点连入
  • 中间节点：每个 joint 节点拆成 X_in / X_out 一对
  • 边 weight：X_in → X_out 的容量 = 保留这个 tensor 的 byte 数（加权后；不存就要重算）

最小割集：所有"被切的边"对应的 tensor 必须保留给 backward。源码在 torch/_functorch/partitioners.py:2167 用 nx.DiGraph() 建图、:2412 调 nx.minimum_cut(nx_graph, "source", "sink") 求解。networkx 没装时（:2045）这套"smart recomputation heuristics"直接报错。

实战上 min-cut 在典型 transformer 规模的图上求解很快，与 trace + functionalize、尤其与后面 Triton 编译的时间相比可忽略。

如果你想看 min-cut 的具体决策，开 TORCH_LOGS=aot_joint_graph,aot_graphs 对照切分前后的图（joint_graph、partitioner 都不是合法的 TORCH_LOGS 名字）；分区器自己的日志（如 §13.20 例子里实跑打出的 materialized backwards: relu (detach, output)）走的是 torch._functorch.partitioners 这个 logger，不在 aot 这个日志名覆盖的模块里，可以用 TORCH_LOGS="+torch._functorch.partitioners" 按模块名打开。这是调试 activation_checkpoint 自动决策的金钥匙。

13.15 AOTAutograd × dynamic shape

第 12 章 §12.5 讲了 Dynamo 的 guards 处理 dynamic shape。AOTAutograd 这一层支持是用 SymInt 在 fake tensor trace 中传递：

下面是示意，不是可运行代码（FakeTensor(shape=...) 不是真实构造方式）：

# Dynamo 看到 batch dim 是 SymInt s0
input.shape = (s0, 768)

# AOTAutograd trace 时 fake input 也是 (s0, 768)
fake_input = FakeTensor(shape=(s0, 768), ...)

# trace 出来的 fx graph 里所有 shape 表达式都是 SymInt
output = forward_graph(fake_input)
# output.shape = (s0, 768)  # 仍是符号

# min-cut partition 按 SymInt 的 hint 估算字节数
# (partitioners.py 的 _size_of: size_hint(numel, fallback=4096))
cost = size_hint(output.numel(), fallback=4096) * dtype_size

代价：cost 估算用的是首次编译时那组具体值的 hint（无 hint 的 unbacked SymInt 退回 4096），partition 决策对其他 shape 未必最优。这种近似够不够好取决于具体模型，源码里没有给出保证。

AOTConfig.dynamic_shapes（§13.8）不是用户直接设的开关：经 Dynamo 进来时，它取决于 fake mode 有没有 ShapeEnv（torch/_functorch/aot_autograd.py:1024-1027）；哪些维度真的是符号，由 Dynamo 决定（torch.compile(dynamic=...) 与 automatic dynamic，第 12 章 §12.5）。

13.16 forward 输出的打包：saved tensors 与 user outputs

trace 完 forward 后，输出包含两类张量：用户实际要的 outputs + backward 需要的 saved tensors。AOTAutograd 把它们一起返回：

# fw_module 的实际输出
fw_outputs = fw_module(*primals)
# = (user_output_0, user_output_1, ..., saved_tensor_0, saved_tensor_1, ...)

记录"前几个是用户输出、后几个是 saved"的不是 AOTConfig，而是 schemas.py 里的 ViewAndMutationMeta（它算出 num_forward_returns / num_forward 等一串下标，见 :644 附近）。runtime_wrappers (§13.7) 按这些下标解包：

# 简化示意：实际 fw 输出依次是 (被改的输入, 用户输出, 中间 base, saved tensors, saved symints)
real_outputs = fw_outputs[:num_user_outputs]   # 给用户
saved = fw_outputs[num_user_outputs:]            # 给 backward

把 saved 与 user output 同一次 forward kernel 计算让二者共享中间状态、避免重复 kernel launch。Inductor 在 codegen 时也按这套语义生成 wrapper。

13.17 AOTAutograd × autocast

第 20 章讲过 autocast。AOTAutograd 怎么处理用户的 with autocast(...): out = model(x)？

机制：Dynamo 把 autocast 区域原样记进 FX 图（本机实跑可见 _enter_autocast('cpu', torch.bfloat16, True, None) / _exit_autocast 两个节点包住 l_a_ @ l_b_）。AOTAutograd 在 dispatcher 层重新 trace 时 autocast 生效，类型转换被展开成显式的 _to_copy 节点，functionalize 与 partition 都在这张 dtype 已经写明的图上做，不需要二次处理。

保存给 backward 的是低精度副本，不会再 cast 回 fp32。本机对 with torch.autocast("cpu", dtype=torch.bfloat16): (a @ b).float().sum() 实跑（min-cut 分区），forward 图先把 a、b _to_copy 成 bf16 再 mm，存给反向的是这两个 bf16 副本的转置（t、t_1）；backward 图把 fp32 的梯度 _to_copy 成 bf16 做两个 mm，再 _to_copy 回 fp32 作为 a、b 的梯度。精度的回升发生在这条 _to_copy 反向链上，而不是保存环节。

13.18 AOTAutograd 与 functorch transform 的协作

第 5 章 §5.5.2 追踪过 functorch 的 vmap(grad(f)) 在 dispatcher 里怎么走。AOTAutograd 与这些变换可以叠加（loss_fn 为任意标量损失函数）：

import torch
from torch.func import vmap, grad

@torch.compile
def per_sample_grad(x_batch, y_batch):
    return vmap(grad(loss_fn))(x_batch, y_batch)

这段代码里：

  1. vmap(grad(...)) 是双重函数变换 —— 每个样本独立算 grad
  2. @torch.compile 编译整段
  3. Dynamo 直接 inline 进 vmap / grad 的 Python 实现，FX 图里记下的是 _vmap_increment_nesting、_add_batch_dim、_wrap_for_grad、_autograd_grad、_remove_batch_dim 这类调用（本机 backend="aot_eager"、TORCH_LOGS=graph_code 实跑），并不是 HigherOrderOperator 节点
  4. AOTAutograd 在 dispatcher 层重新 trace 这张图时经过 functorch 的变换层，最终编译产物里 vmap 被展开成 batched op，grad 被展开成 backward 链

per-sample gradient（用于 differential privacy 训练）通过这条路径拿到的是"批量化 + 融合"后的 kernel，而不是逐样本 Python 循环，量级差距可以很大（具体倍数取决于 batch 与模型，本机无 GPU，不给数）。AOTAutograd 把"复杂的函数变换"变成"普通可编译 op"，让用户能用 JAX 风格 API + PyTorch eager 体验 + Inductor 性能。

13.19 性能数字：AOTAutograd 在编译时间里的占比

完整 torch.compile 链路的编译时间按段拆开大致是这个次序（本机无 GPU，无法给实测秒数；torch._dynamo.utils.compile_times() 会按 phase 打印你自己机器上的真实数字）：

阶段 相对量级
Dynamo trace 小
AOTAutograd trace + functionalize 中
AOTAutograd partition（min-cut） 小
Inductor lowering + scheduling 中
Triton kernel 编译（多进程并发） 最大头

AOTAutograd 在编译时间里占据可观的一段。这是它在编译栈里的"位置成本" —— 不是最重的（Triton 编译最贵），但显著。跨进程重启时能跳过这一段的是 AOTAutogradCache（torch/_functorch/_aot_autograd/autograd_cache.py，torch._functorch.config.enable_autograd_cache 默认开启），它与缓存 Inductor 产物的 FxGraphCache（§14.11.2）配合使用；Dynamo 的 trace 仍要重做。

13.20 一个完整 trace 例子：F.linear(x, w, b).relu()

把整章串起来看一段简单代码的 AOTAutograd trace：

@torch.compile
def f(x, w, b):
    return F.linear(x, w, b).relu()

AOTAutograd 收到 fx graph 后：

下面各步的图取自本机实跑：backend="aot_eager"、TORCH_LOGS=aot_joint_graph,aot_graphs，x 为 [4, 8] 且不需要 grad，w 为 [16, 8]、b 为 [16]，二者需要 grad（去掉了类型注释与 = None 释放语句）。

  1. trace joint（forward + backward）：
   def forward(self, primals, tangents):   # primals_1/2/3 = x, w, b; tangents_1 = 输出梯度
       t = aten.t(primals_2)
       addmm = aten.addmm(primals_3, primals_1, t)             # linear
       relu = aten.relu(addmm)
       detach = aten.detach(relu)
       detach_1 = aten.detach(detach)
       threshold_backward = aten.threshold_backward(tangents_1, detach_1, 0)   # relu 反向
       t_1 = aten.t(threshold_backward)
       mm = aten.mm(t_1, primals_1)                            # w 的梯度
       t_2 = aten.t(mm)
       sum_1 = aten.sum.dim_IntList(threshold_backward, [0], True)   # b 的梯度
       view = aten.view(sum_1, [16])
       t_3 = aten.t(t_2)
       return [relu, None, t_3, view]                          # x 不要梯度, 对应 None

  1. functionalize：原始 graph 没 inplace，没有需要改写的节点

  2. partition (min-cut)：哪些中间值保留给 backward？

    • addmm（linear 输出）：不存。relu 的反向只需要 relu 的输出
    • relu 的输出（以 detach_1 的形式）：保留，threshold_backward 要用
    • primals_1（即 x）：保留，算 w 的梯度要用
    • 决策：保存 primals_1 + detach_1

  3. 拆出 fw / bw：

   fw_module: (primals_1, primals_2, primals_3) → (relu, primals_1, detach_1)
   bw_module: (primals_1, detach_1, tangents_1) → (None, t_3, view)
  1. 送 Inductor：fw_module 与 bw_module 分别交给 Inductor 编译（本机无法跑 Inductor 代码生成，具体切成几个 kernel 未实测）

整套流程一次就跑完（此后命中缓存）。用户写 3 行 Python，AOTAutograd 自动展开成正反向 trace + 分片 + 编译，是 PyTorch 训练加速最复杂的中间环节。

13.21 trace 阶段的 fake_mode 详解

§13.3 提了"用 fake tensor 跑 autograd.grad"。具体怎么实现？

FakeTensorMode 是个 TorchDispatchMode（第 5 章 §5.7）。进入 mode 后所有算子调用被拦截：

import torch
from torch._subclasses.fake_tensor import FakeTensorMode

real_x = torch.randn(3, 4)
with FakeTensorMode() as fake_mode:
    fake_x = fake_mode.from_tensor(real_x)    # 把真实 tensor 转成 fake
    fake_y = fake_x + 1                         # 算子被拦截, 只算 shape, 不真做加
    print(fake_y.shape)                          # torch.Size([3, 4]); 但没有真实数值, .tolist() 会抛异常

AOTAutograd trace 时全程在 FakeTensorMode 里。fake_y 看起来像 tensor 但只有 shape / dtype / device 元信息 —— 没分配显存、没真做计算。这让 trace 70B model 在小显存机器上也能跑（trace 不需要装下完整模型）。

autograd.grad(fake_loss, fake_inputs) 在 fake tensor 上跑，触发反向链 trace。反向规则的"中间值"也是 fake，最终 fx graph 上每个节点是一个 fake op。

整套 fake trace 让"先 trace 出 graph、再用真实 tensor 编译运行"成为可能。这是 AOTAutograd 能处理超大模型的根本。

13.22 flat_args 与 FlatArgsAdapter：嵌套结构展平

用户的 forward 输入可能是嵌套结构（dict / list / tuple）：

def model(x, kwargs={'attn_mask': mask, 'cache': (k, v)}):
    ...

AOTAutograd / Inductor 内部处理纯 tensor 列表 —— 嵌套结构必须先 flatten。这件事由 torch/utils/_pytree.py 的 tree_flatten / tree_unflatten 完成。（另外，export 路径上还另配了一个适配器 FlatArgsAdapter（torch/export/unflatten.py:303），用来把调用方的实参对齐到导出时记录的 TreeSpec。）

下面是示意，不是 PyTorch 源码原文——fw_module 是为讲清 flatten / unflatten 关系现编的名字，TreeSpec(...) 也只是占位。

要看的是这个形状：编译路径拿到的永远是扁平 tensor 列表，嵌套结构靠 spec 在两端还原。

flat_args = [x, mask, k, v]                 # 扁平 tensor 列表
spec = TreeSpec(...)                          # 描述如何重建嵌套

# trace 时 forward 接 flat_args, 内部用 spec unflatten 还原嵌套
def fw_module(*flat_args):
    x, mask, k, v = flat_args
    kwargs = {'attn_mask': mask, 'cache': (k, v)}
    return model(x, **kwargs)

pytree（PyTorch 的嵌套结构工具，torch/utils/_pytree.py）提供 flatten / unflatten 实现。这套机制让"任意嵌套 input"都能进编译路径，不需要用户改 forward 签名。

13.23 synthetic_base：处理 view 别名

AOTSyntheticBaseWrapper（§13.7 提过）处理一个微妙场景：多个 input 张量是同一个 base 的不同 view：

def f(x, y):
    x.mul_(2)          # 原地改 x
    return x + y

base = torch.randn(10)
f(base[:6], base[4:])     # 两个 input 共享 storage, 且 base[4:6] 重叠

如果 trace 时把 x / y 当独立 tensor，functionalize 后两者的 mutation 不会互相影响 —— 但真实运行时它们共享 storage、修改 x 会影响 y 的重叠部分。编译产物在这种场景下行为错误。

AOTSyntheticBaseWrapper 检测这种共享 base 场景，把多个 view input 合并成"传一个 base"。trace 时 base 是 leaf input，view 在 graph 内部用 as_strided 重新构造，写回用 as_strided_scatter —— functionalize 与 mutation 分析在 base 层面正确进行。本机对上面的例子调 aot_function 实跑，编译器拿到的只有一个 [10] 的输入，图里是 as_strided → mul → as_strided_scatter → 两个 as_strided → add。

这套机制让"用户传 view 给 compiled function" 在边角情况下仍正确。它没有开关：AOTSyntheticBaseWrapper.pre_compile 先跑 merge_view_inputs()，返回 synthetic_base_info is None 时走 happy path 直接返回，否则一定合成 base。merge_view_inputs() 返回 None 的情形包括：没有任何输入被原地改数据、互为别名的输入里没有被改的、以及共享 storage 但实际内存不重叠（例如把上例换成 f(base[:5], base[5:])，本机实跑两个输入原样保留）。反过来，两条路径被显式禁掉并直接抛 RuntimeError：输入里有 tensor subclass（requires_subclass_dispatch）、以及 aot_config.is_export 的导出路径（torch/_functorch/_aot_autograd/runtime_wrappers.py:1385-1400）。

13.24 AOTDedupeWrapper：相同张量去重

另一个 corner case：用户把同一个张量作为多个 input 传：

def f(x, y):
    x.add_(1)
    return x + y

result = f(t, t)    # x 和 y 是同一对象, 且被原地改

trace 时 graph 里会同时引用 t 两次。AOTDedupeWrapper（torch/_functorch/_aot_autograd/runtime_wrappers.py:1106）检测这种情况。它分两种策略：重复的输入没有被原地修改时，只把第二份换成一个 detach 出来的独立叶子张量，两个形参都保留（本机把上例去掉 x.add_(1) 实跑，图仍是两个输入）；只有重复输入被原地改了，才真正去重，trace 时把 t 只列一次输入。上例本机调 aot_function 实跑的图（默认 keep_inference_input_mutations=False，所以 mutation 的新值作为第一个输出返回）：

def forward(self, arg0_1):                          # 只一个 input
    add = torch.ops.aten.add.Tensor(arg0_1, 1);  arg0_1 = None
    add_1 = torch.ops.aten.add.Tensor(add, add)     # 同一 input 用两次
    return (add, add_1)

去重让 mutation 在别名输入之间的传播在图里是正确的；运行时 wrapper 再按 keep_arg_mask 把传进来的重复参数剔掉，再调编译产物。

13.25 RNG 处理：让 functionalize 与重算一致

§13.9 的 banned ops 列表里有 aten.native_dropout 等随机算子。但这只解决了"min-cut 不主动重算随机算子" —— 真正的 RNG 处理更复杂。

问题：activation_checkpoint 在反向时重 forward 取 activation。如果 forward 包含 dropout，第二次重 forward 时 RNG state 不同，得到的 dropout mask 也不同 —— 梯度计算错乱。

AOTAutograd 分两种情况处理：

  • 默认情况：随机算子不在可重算白名单里（§13.9），partitioner 根本不重算它，而是把结果存下来。本机对 F.dropout(x.sin(), p=0.5, training=True).cos().sum() 用 aot_eager 实跑，forward 图保存的是 native_dropout 产出的 bool mask（getitem_1: "b8[8]"），backward 直接拿它调 native_dropout_backward
  • 用户 checkpoint 标了要重算的随机算子：RNG state 也作为 saved tensor。torch/_functorch/partitioners.py:1607 的 functionalize_rng_ops 把 forward 里的随机算子换成 run_and_save_rng_state，把 RNG state 作为输出存下；backward 里的重算版换成 run_with_rng_state，先恢复 RNG state，保证两次 forward 的随机数一致

torch/_inductor/fx_passes/replace_random.py（§14.7.1 的 Pass）里的 replace_random（:99）把 aten.rand / aten.randn 替换成 inductor_prims.random(size, inductor_prims.seed(device), mode)，让 seed 显式变成一个图节点。它属于 Inductor 的 joint graph pass，在 partition 之前运行。

理解这条机制让你看到 torch.compile 下的 dropout 不会因为 checkpoint 重算引入 forward / backward 之间的随机数不一致。注意这不等于与 eager 逐位相同：Inductor 默认用自己的随机数生成（torch._inductor.config.fallback_random = False），要与 eager 的随机数流一致需把它设成 True。

13.26 AOTAutograd 历史演进：从 functorch 到主仓

AOTAutograd 不是从零写的，是 functorch 项目（2022 年发布）的核心组件之一。functorch 借鉴 JAX 的函数变换思想（vmap / grad / jvp），把 PyTorch 的 autograd 改造成"函数式" —— 让"对函数求导"成为编译期可处理的对象。

时间线（括号里是各 tag 所指提交的日期，不是 GitHub Release 发布日）：

  • v1.13 (2022-10)：functorch 源码并入主仓（functorch/ 目录），随 PyTorch 一起发布，含 AOTAutograd 雏形
  • v2.0.0 (2023-03-09)：functorch 的实现迁入 torch/_functorch（torch/_functorch/aot_autograd.py 首次出现于 v2.0.0），AOTAutograd 成为 torch.compile 的核心组件
  • v2.1.0 (2023-09-21)：Compiled Autograd 的实现（torch/_dynamo/compiled_autograd.py）首次随 release 发布
  • v2.4.0 (2024-07-09)：AOTAutograd 与 DTensor / FSDP-2 / export 深度集成
  • v2.11.0 (2026-03-20)：本专栏基准版本

理解这条演进让你看 torch._functorch 模块的命名时不困惑：functorch 这个老名字保留下来作为 namespace，但实际功能远超原始 functorch（vmap / grad）—— 它是整个 AOT 编译框架的代号。

13.27 AOTAutograd × Compiled Autograd

Compiled Autograd 把反向 Engine（第 8 章）也 trace 进编译路径。普通 AOTAutograd 把 backward 拆成 bw_module 编译，但调用 bw_module 仍由 Engine 调度。Compiled Autograd 让 Engine 调度本身被 Dynamo trace、Inductor 编译。

import torch._dynamo
torch._dynamo.config.compiled_autograd = True

# 之后 loss.backward() 的整个调度也被编译

收益：消除 Engine 逐节点调度的开销，并让整条反向链有机会被 Inductor 融合。反向本身 kernel 越碎、Engine 调度占比越高，这块收益越大。

实现机制：torch/_dynamo/compiled_autograd.py 的 AutogradCompilerInstance（:285）配合 C++ 侧的 torch/csrc/dynamo/python_compiled_autograd.cpp，在 Engine 走 grad_fn 链时把每个节点的 apply 记录成 fx graph，而不是当场执行。开关是 torch._dynamo.config.compiled_autograd（torch/_dynamo/config.py:719，默认 False）。

这是 PyTorch 编译器栈的"最后一块拼图" —— 之前 forward / backward / optimizer step 都能编译，现在 Engine 调度也加入。理解了 §13 + §8 + §14，你就理解了为什么 Compiled Autograd 是自然演进的下一步。

13.28 AOTAutograd 错误诊断速查表

症状 可能原因 诊断
trace 自定义算子时报缺 fake / meta 实现 fake tensor 上跑不动该算子 自定义算子缺 register_fake
backward 数值与 eager 不一致 RNG 处理出错 检查 dropout / random op 是否被正确 functionalize；逐位对比时先开 fallback_random
编译后 OOM 但 eager 不会 min-cut 决策不优 调小 unit 粒度或显式 activation_checkpoint
图里出现 as_strided_scatter、输入个数比预期少 synthetic_base 触发（重叠的 view 输入被原地改） 改用 .clone() 让 input 独立
AOTAutograd trace 时间长 model 太大 + min-cut 复杂 靠 AOTAutogradCache + FxGraphCache 缓存（§13.19、§14.11.2）

把这套表配合 TORCH_LOGS=aot,aot_graphs,aot_joint_graph 用，能快速定位 AOTAutograd 阶段的问题。

13.29 AOTAutograd 中的"trace 的 trace"

AOTAutograd 的 trace 阶段实际上调用了 PyTorch 的 dispatcher trace 机制（第 5 章 §5.7 TorchDispatchMode）—— 用 ProxyMode 把每次算子调用记录成 fx Node。这与 Dynamo trace 是两个不同层面：

  • Dynamo trace：在 Python 字节码层，把 Python 函数翻译成 fx Graph
  • AOTAutograd trace：在 dispatcher 层，把 fx Graph 上每个 torch 调用通过 fake tensor + ProxyMode 展开、重新 trace 成 ATen 级的 fx Graph（含反向）

所以 AOTAutograd 是 "trace 的 trace" —— Dynamo 的输出是 AOTAutograd 的输入，AOTAutograd 再往下 trace 一层。这种"多层 trace 串联"是 PyTorch 编译器栈的核心设计。

13.30 几个 corner case

AOTAutograd 复杂度的来源是无数 corner case：

  • Tensor subclass：FakeTensor、TwoTensor、自定义 subclass 的 trace 路径
  • View ops：x.view(...) 在 functionalize 时要重新计算 view 关系
  • AMP autocast：joint graph 要保留 amp 上下文，让 backward 用同样的 dtype 策略
  • Checkpoint interaction：用户已经在用 torch.utils.checkpoint 时，AOTAutograd 要识别并尊重 checkpoint 边界
  • Inplace foreach：optimizer 的 _foreach_* inplace 集合操作要特殊 functionalize
  • Side-effect ops：rng / print / collective ops（NCCL）不能被随便 reorder

每个 corner case 在源码里都有专门的 if 分支。这是 AOTAutograd 代码量的真正来源 —— 绝大多数代码不在描述编译算法，而在处理边界情况。

13.31 effect tokens：side effect 算子的顺序保留

某些算子有 side effect：aten::_print、torchbind 对象的方法调用（call_torchbind）等，在 torch/_higher_order_ops/effects.py:63-67 里被注册为 effectful op（自定义算子也可以通过 torch.library._register_effectful_op 登记）。functionalize 不能简单消除它们，否则程序行为变化。

EffectTokensWrapper（§13.7）的解法：把 side effect 表达成 token 依赖链：

op_with_effect_1 → token_1
                    ↓
op_with_effect_2 → token_2
                    ↓
op_with_effect_3 → token_3

每个 side effect op 接收 input token、产出 output token。后续 op 必须等前面的 token —— 形成显式依赖链。这套机制让 functionalize 能"假装这些 op 是纯函数"（只是依赖 token），但实际执行顺序与原始一致。

实战提醒：collective ops（如 all_reduce）不走 effect token——它们不在上面的 effectful 注册表里。functional collectives（第 16 章 §16.12）返回新张量、再用 wait_tensor 取结果，顺序靠数据依赖维持；partitioner 默认也不重算 _c10d_functional 命名空间的算子（torch/_functorch/partitioners.py:2174 附近）。§14.30 的 reorder_compute_for_overlap 那组 reorder pass 调整的正是这类 collective 与计算的相对位置。

13.32 fx graph 上的 CSE 与简化

trace 完成后 fx graph 可能含大量冗余（同一个表达式重复算）。送给 Inductor 前，joint graph 会经过几个简化 pass，其中前两个是 AOTAutograd 分区器自己做的，后两个来自 Inductor：

  • CSE (Common Subexpression Elimination)：相同表达式只算一次（与第 14 章 §14.9.3 的 codegen 层 CSE 类不同 —— 这是 fx 层的、那是生成 kernel 代码时的）
  • dead code elimination：没用到的中间值删掉
  • constant folding：编译期能算的常量提前算
  • algebraic simplification：x * 1 → x、x + 0 → x 等

这些 pass 减小 graph 大小，让后续 Inductor 工作量降低。fx 层的 CSE 实现是 torch/_functorch/compile_utils.py:51 的 fx_graph_cse，由 min_cut_rematerialization_partition 在开头按 torch._functorch.config.cse（默认 True）调用；常量折叠（constant_fold_uniform_value）与 remove_no_ops 这类化简在 torch/_inductor/fx_passes/joint_graph.py，由 Inductor 交给 AOTAutograd 的 partition_fn 在切分前调用。

CSE 在 AOTAutograd 这层做的好处是：functionalize 后会引入大量重复（如多个 view 都涉及 base + offset 计算），CSE 立即消除让送给 Inductor 的图清爽。

13.33 AOTAutograd 与 Hooks 的协作

第 7 章 §7.5.3 / 第 9 章 §9.8 提过 hooks（saved_tensors_hooks / forward_hook / backward_hook 等）。AOTAutograd 编译路径下这些 hook 怎么处理？

  • saved_tensors_hooks：v2.11 的 AOTAutograd 在 hook 可内联时，把 pack hook 接到 forward 图末尾、unpack hook 接到 backward 图开头（torch/_functorch/_aot_autograd/graph_compile.py:1179 的 maybe_inline_graph_saved_tensors_hooks；源码注释说主要用例是激活量化），默认只作用于图内部的中间值（saved_tensors_hooks_filtering_mode = "donated"）
  • forward_pre_hook / forward_hook：trace 时 Dynamo 把 hook 调用 inline 进 graph（被当成普通函数调用 trace）—— 编译产物里 hook 自动展开；backward 侧的 hook 要到反向运行时才触发，不在这张 forward 图里
  • register_post_accumulate_grad_hook：FSDP / DDP 用，AOTAutograd 不直接支持 —— 这些 hook 在 graph 之外触发，与编译路径协作通过 collective op 与 functional collectives（第 16 章 §16.12）

理解 hook 在编译路径下的处理让你看到"用户写的 hook 在 compile 后是否还生效"。多数 hook 仍生效（trace 时被 inline）；少数与训练循环协作的 hook（如 DDP）通过专门机制保留。

13.34 一段 LLM Attention 的 AOTAutograd trace 简要

把 §13.20 的简单例子升级到 LLM attention：

@torch.compile
def attention(q, k, v, mask):
    # q, k, v: [B, H, S, D]
    scores = q @ k.transpose(-2, -1) / sqrt(D)    # [B, H, S, S]
    scores = scores + mask
    attn = F.softmax(scores, dim=-1)
    out = attn @ v                                   # [B, H, S, D]
    return out

AOTAutograd 处理：

  1. trace joint：捕获 forward + backward。backward 中 attention 反向涉及 softmax 反向（complex jacobian）
  2. functionalize：这段代码没有 inplace（scores = scores + mask 是普通加法），没有需要改写的节点
  3. Inductor 的 joint graph pass：在 partition 之前，torch/_inductor/fx_passes/fuse_attention.py 里的模式（_sfdp_pattern_*）若匹配上，整段 matmul → scale → softmax → matmul 被替换为 aten.scaled_dot_product_attention，之后再按设备与 shape 分派到 flash / memory-efficient 等实现
  4. partition (min-cut)：没匹配上模式时，softmax 输出 attn 这类中间值是否保留由 min-cut 决定 —— [B, H, S, S] 的 attention 概率矩阵是 LLM 训练 activation 的主要内存压力来源之一

模式匹配成功时，编译产物里 attention 整段交给 SDPA 的融合实现，反向也走对应的反向 kernel。用户写朴素 attention，只要写法命中 fuse_attention 里登记的模式，就能用上融合 attention —— 这是 AOTAutograd + Inductor 配合的典型工程价值（本机无 GPU，未实跑）。

13.35 AOTAutograd 与第三方 backend 的协作

§12.7 提过 aot_eager / cudagraphs / inductor 三种 backend。AOTAutograd 是这些 backend 的"上游" —— trace 完后把 graph 给 backend：

# 用户调 torch.compile(model, backend='inductor')
# Dynamo 输出 forward graph → AOTAutograd trace + partition →
# 把 fw_module 与 bw_module 分别送给 Inductor 的 compile_fx_inner (torch/_inductor/compile_fx.py:787)

第三方 backend 接入流程。注意：直接用 @torch._dynamo.register_backend 注册的函数拿到的是 Dynamo 输出的 torch 级 FX 图，还没经过 AOTAutograd（本机实跑，gm.code 里是 l_x_ @ l_w_、matmul.relu()）；要拿到切分后的 ATen 级 fw / bw 图，需用 torch._dynamo.backends.common.aot_autograd 把自家编译器包一层：

import torch
from torch._dynamo.backends.common import aot_autograd
from functorch.compile import make_boxed_func

def my_compiler(gm, example_inputs):
    # gm 是 AOTAutograd 切分后的 fw_module 或 bw_module（ATen 级）
    return make_boxed_func(gm.forward)    # 换成自家编译逻辑

my_backend = aot_autograd(fw_compiler=my_compiler, bw_compiler=my_compiler)
model = torch.compile(model, backend=my_backend)

这套接入路径让国内厂商能给 PyTorch 加自家 backend，不需要修改主仓代码。AOTAutograd 给所有 backend 提供"标准化的可编译 IR"（functionalize + partition 后的 fx graph），是 PyTorch 编译器生态的工程基石。

13.36 AOTAutograd × Higher Order Operator (HOP)

第 14 章 §14.11.4 提过 HigherOrderOperator —— 算子的"参数"是另一段计算图（如 cond / while_loop / flex_attention）。AOTAutograd 处理 HOP 时有特殊逻辑：

  • 在主图里保留为单个节点：HOP 调用不被展开，子图作为参数挂在这个节点上
  • 每个 HOP 自带各层规则：HOP 为 Functionalize、ProxyTorchDispatchMode、Autograd 等分别注册实现（如 torch/_higher_order_ops/cond.py:689 的 @cond_op.py_functionalize_impl），AOTAutograd trace 到它时，由这些规则递归地把子图函数化、trace 成 ATen 子图
  • 反向子图由 HOP 自己生成：如 cond 的 CondAutogradOp（torch/_higher_order_ops/cond.py:295）用 create_bw_fn 为两个分支分别构造反向子图

HOP 让 PyTorch 能 trace "动态控制流"代码（写成 torch.cond(pred, true_fn, false_fn, operands) 而不是 Python 的 if），同时让每个分支被独立编译。这是与 JAX 的 lax.cond / lax.while_loop 思想一致 —— 控制流作为 first-class IR 节点。

实战：FlexAttention（自定义 attention pattern）就是用 HOP 实现的。用户写 score_mod 函数描述自定义 attention 行为，HOP 把它 trace 成子图、AOTAutograd 给它产出反向、Inductor 编成 Triton kernel。整套机制让自定义 attention 享受与内置 SDPA 同等的编译优化。

13.37 AOTAutograd 与 export 的特殊处理

torch.export 与 AOTAutograd 的关系在 v2.11 分两步：torch.export.export() 默认 strict=False、pre_dispatch=True，先用 make_fx 产出训练 IR（torch/export/_trace.py 的 _export_to_aten_ir_make_fx）；等调用 ep.run_decompositions() 降到 ATen IR 时，才经 torch/export/exported_program.py 调 _export_to_aten_ir，走 AOTAutograd 的 aot_stage2_export（§13.11）。这条 export 路径与 autograd / inference 路径的关键差异：

  • 不允许 graph break：export 要求"严格静态 graph"，遇到 unsupported op 直接报错，不能 fallback eager
  • 不允许 dynamic shape（除非用户显式 dynamic_shapes={...} 标记）
  • 不 trace backward：export 的目标是部署，不是训练
  • 保留更多语义信息：export 输出的 ExportedProgram 含 module hierarchy、原始 Python 调用栈，比纯 fx graph 信息丰富

ExportedProgram 是 v2.x 部署的核心 IR。AOTI（第 15 章 §15.10.1）从 ExportedProgram 出发编译，ExecuTorch 等移动端部署方案也以 export 为入口。理解 AOTAutograd 在 export 路径上的角色让你能跟踪"训练时的 model 怎么变成部署时的 .so"完整链路。

13.38 AOTAutograd 的并发与多线程

AOTAutograd trace 本身是单线程（CPython GIL 限制），但它产出的 graph 编译阶段（Inductor）是多进程的（§14.11.1 AsyncCompile）。整体编译时间分布：

  • AOTAutograd trace + functionalize：单线程
  • AOTAutograd partition：单线程
  • Inductor lowering + scheduling：单线程
  • Triton 编译：多进程并发

各段在你机器上的实际耗时，用 torch._dynamo.utils.compile_times() 看（本机无 GPU，不给秒数）。

整体编译时间瓶颈在 Triton（§13.19 表格）。AOTAutograd 的单线程性质不是问题 —— trace 阶段比 Triton compile 快得多。

未来方向：AOTAutograd trace 可能也并行化（用 thread pool 同时 trace 不同 subgraph），但实现复杂度高、收益相对小。当前设计够用。

13.39 AOTAutograd 在 PyTorch 演进路线上的位置

AOTAutograd 在 PyTorch 整体演进中的位置：

graph LR
    Eager[v1.x eager]
    JIT[v1.x TorchScript]
    Functorch[v1.13 functorch<br/>函数变换]
    Compile[v2.0 torch.compile<br/>= Dynamo + AOTAutograd + Inductor]
    Compiled[Compiled Autograd<br/>反向也进编译]
    Future[v3.x ?<br/>更激进编译 / DTensor 全自动并行]
    Eager --> JIT
    Eager --> Functorch
    Functorch --> Compile
    Compile --> Compiled
    Compiled --> Future
    style Compile fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

AOTAutograd 是 v2.0 编译栈的核心 —— 没有它就没有 torch.compile 的工程实现。理解 AOTAutograd 让你能解读 PyTorch 团队公开的 RFC、跟上未来演进路线。

13.40 几条工程经验

实战 AOTAutograd 相关：

1. 调试时用 aot_eager backend：torch.compile(fn, backend='aot_eager') 跳过 Inductor，只跑 AOTAutograd。如果这步 OK 但 Inductor 阶段错，就能快速定位

2. TORCH_LOGS=aot_joint_graph,aot_graphs 看到 trace 出来的 joint graph + 切分结果：知道哪些张量被保留、哪些被重算

3. 有的算子 functionalize 失败：通常是用户写的自定义算子缺 register_fake。第 22 章自定义算子会展开

4. 用 torch._functorch.config 里的 partition 开关调偏好：activation_memory_budget（默认 1.0 = 默认的 runtime-optimized 策略，0.0 = 对整个编译区做 activation checkpointing，中间值如 0.4 表示只存默认策略 40% 的 activation；内部解一个 0-1 背包）、activation_memory_budget_solver（默认 "dp"，另有 "greedy" / "ilp" / "dp_knapsack_sliding_hirschberg"）、recompute_views（默认 False）

5. 与 torch.utils.checkpoint 的协作：被 checkpoint 区段覆盖的节点带有 CheckpointPolicy 标记，partitioner 对 MUST_RECOMPUTE / PREFER_RECOMPUTE 强制重算、对 MUST_SAVE 强制保存（§13.9）；区段之外的节点仍由 min-cut 决定

13.41 AOTAutograd 在用户代码里的"可见性"

普通用户写 @torch.compile 时不知道 AOTAutograd 的存在 —— 它对用户透明。但有几个 API 让用户能间接感知 / 控制：

  • torch.export.export(model, ...) 在 v2.11 默认不经 Dynamo（strict=False），ep.run_decompositions() 时走 AOTAutograd 的 aot_stage2_export（§13.37）
  • torch._functorch.aot_autograd.aot_function(fn, fw_compiler, bw_compiler)（也可 from functorch.compile import aot_function）直接暴露 AOTAutograd 接口（不经 Dynamo）—— 用于研究 / 调试，本章几处实跑就是用它
  • AOTConfig 的多数字段由 torch._functorch.config / torch._inductor.config 间接决定（如 torch._inductor.config.freezing 会影响 inference 路径的常量折叠）

这些 API 让"想直接玩 AOTAutograd"的高级用户有入口。生产代码用 @torch.compile 即可，不需要直接接触 AOTAutograd —— 这本身就是它的设计目标：对常规用户透明、对高级用户可控。

13.42 跨专栏关联

  • 《Rust 编译器与运行时揭秘》专栏第 15 章 MIR 优化：AOTAutograd 的 functionalize → partition → optimize 与编译器里 LLVM mem2reg 把内存读写提升为 SSA 值、再做 CFG 化简的思路相近 —— 先把 mutation 化简、再做图变换（mem2reg 是 LLVM 的 pass；rustc 的 MIR 本身不是 SSA 形式）
  • 《Serde 元编程》专栏第 12 章 Serialize codegen：functionalization 让算子变纯函数，与 Serde 的"派生宏生成纯函数式 serde"思路相通
  • 第 7 章 §7.5.3 saved_tensors_hooks：AOTAutograd 的 saved tensors 与 eager autograd 的 SavedVariable 是相同概念在两个执行模式下的对应

13.43 AOTAutograd 的"3 段式"工作流总览

把整章串起来，AOTAutograd 内部是清晰的"3 段式"流程：

flowchart TB
    Input[Dynamo 输出: forward FX Graph + Guards]
    Input --> S1[段 1: trace<br/>用 fake tensor + autograd.grad<br/>产出 joint forward+backward graph]
    S1 --> F1[FunctionalTensorMode 拦截 inplace]
    F1 --> S2[段 2: rewrite<br/>functionalize + decomp + dedup<br/>产出纯函数式 IR]
    S2 --> S3[段 3: partition<br/>min-cut 切 fw / bw<br/>用 SymInt 估算 cost]
    S3 --> O[输出: fw_module + bw_module + saved_indices]
    style S1 fill:#fef3c7
    style S2 fill:#dbeafe
    style S3 fill:#dcfce7

每段对应一组源码文件：

  • 段 1：torch/_functorch/aot_autograd.py 的 create_aot_state + torch/_functorch/_aot_autograd/graph_capture.py 的 trace 函数
  • 段 2：torch/_subclasses/functional_tensor.py + torch/_decomp/decompositions.py + torch/_functorch/_aot_autograd/runtime_wrappers.py 的 dedupe / synthetic_base
  • 段 3：torch/_functorch/partitioners.py 的 min-cut 算法

这种"段-段-段"工作流让 AOTAutograd 的代码组织清晰、错误能精确定位到段。

13.44 设计层面的 AOTAutograd vs JAX

AOTAutograd 的思想直接受 JAX 影响。对比：

维度 JAX AOTAutograd
trace 时机 jit 时 torch.compile 第一次调用时
autograd 风格 函数变换 (grad(f) 返回新函数) 命令式 (loss.backward())
inplace 处理 不允许 inplace functionalize 后允许
输入 shape trace 时绑定（dynamic 用 abstract shape） SymInt 支持 dynamic
编译后端 XLA Inductor

JAX 的"函数式"哲学让 AOTAutograd 设计成立 —— 没有 functionalize 这套机制，PyTorch 大量 inplace / view 的代码很难交给编译器做图优化。AOTAutograd 在保留 PyTorch 命令式 API 的前提下，用 functionalize 做"内部函数式"，是连接 eager / compile 两种范式的关键工程胶水。

理解这条 JAX → AOTAutograd 演进让你能预判：未来 PyTorch 若引入更激进的"函数式优化"（如自动 vmap / grad 组合），AOTAutograd 是天然落地点。

13.45 AOTAutograd 在 LLM 训练里的实际收益

编译栈各层的收益来源是可以分开归因的（本机无 GPU，下面只给方向，不给数字；要量就逐层换 backend 跑 eager → aot_eager → inductor → inductor + mode="reduce-overhead"）：

配置 主要收益来源
eager 基线
backend="eager"（只 Dynamo） 基本无收益：调试 backend，原样运行 Dynamo 抓到的图，用来排查 Dynamo 本身
backend="aot_eager"（+ AOTAutograd partition） min-cut 自动 activation checkpoint，降显存峰值
backend="inductor" 算子融合 + Triton codegen
mode="reduce-overhead"（+ CUDA Graph） 消除 kernel launch overhead

AOTAutograd 这一层省下的主要不是 compute（重算反而多）—— 是显存峰值，让用户能开更大 batch_size、训练吞吐间接提升。

理解 AOTAutograd 的工程价值不是"它直接让训练快多少"，是 它让"自动决定 activation_checkpoint 粒度"成为可能。手动 checkpoint 调到最优费时又容易错；AOTAutograd 的默认决策无需调参，需要更省显存时再调 activation_memory_budget（§13.40）。

13.46 AOTAutograd × profiler

第 21 章讲过 PyTorch profiler。AOTAutograd 编译产物在 profiler 里有特殊处理：

  • 编译产物的 fw / bw 在 trace 里各显示为一个事件：CompiledFunction 与 CompiledFunctionBackward，外面还包着 Torch-Compiled Region: 0/0（本机 CPU、aot_eager 实跑）
  • 内部 Triton kernel 显示为 triton_per_fused_xxx（第 14 章 §14.25）
  • AOTAutograd 自身的 trace 阶段（编译时）不在训练 trace 里（编译完成后才训练）

如果你的训练慢，而 profiler 里在稳态 step 中还能看到编译相关的耗时，多半是在反复重编译：先用 TORCH_LOGS=recompiles 查原因；跨进程重启后的编译时间则靠 AOTAutogradCache / FxGraphCache（TORCHINDUCTOR_CACHE_DIR=... 指定缓存目录）节省。

13.47 AOTAutograd 不能做的事

明确 AOTAutograd 的边界：

  • 不处理纯 Python 控制流：那是 Dynamo 的职责，AOTAutograd 拿到的已经是确定的 fx graph
  • 不直接生成 kernel：那是 Inductor 的职责，AOTAutograd 输出的是 fx graph，不是 Triton 代码
  • 不优化单算子内部：算子级优化（如 fp16 GEMM 是用 cuBLAS 还是 Triton template）由 Inductor 决定
  • 不处理跨进程通信：那是 ProcessGroup 与 functional collectives 的职责（第 16 章），AOTAutograd 把 collective 当作普通算子（不带 effect token，顺序靠数据依赖，§13.31）

明确边界让你不会"在 AOTAutograd 找其他模块的功能"。当编译某段代码出错时，先判断是哪个模块的问题：Dynamo（trace 失败 / graph break）、AOTAutograd（functionalize / partition 错）、还是 Inductor（codegen 失败）。

13.48 AOTAutograd × FSDP-2 的协作

第 18 章 §18.19 讲过：v2.11 默认 skip_fsdp_hooks=True，FSDP-2 的 hook 在图外 eager 执行，AOTAutograd 只看到计算。下面说的是把它设为 False 的 Traceable FSDP2 模式，具体到 AOTAutograd 这层：

  • DTensor 走 subclass desugaring：DTensor 是可 trace 的 wrapper subclass，AOTAutograd 通过 __tensor_flatten__ 把它拆成内部的普通张量来 trace，图里只剩普通张量；placement 等元信息记在 SubclassMeta 里，由 AOTDispatchSubclassWrapper（§13.7）在运行时重新包回 DTensor
  • collective 以 functional 形式进图：FSDP-2 源码里调的是传统的 dist.all_gather_into_tensor / dist.reduce_scatter_tensor（torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py），Dynamo trace 时按 torch/distributed/_functional_collectives.py 的 traceable_collective_remaps 把它们换成 functional collectives（第 16 章 §16.12），AOTAutograd 再把它们当作普通算子加进 graph
  • collective 不被重算：partitioner 默认不重算 _c10d_functional 算子（§13.9），执行顺序靠 wait_tensor 的数据依赖维持，而不是 effect token

这套机制让 FSDP-2 训练 graph 能完整经过 AOTAutograd → Inductor 编译。FSDP-1 的 FlatParameter 在 AOTAutograd 路径上引发各种特殊处理，是它与 compile 兼容性差的根本。FSDP-2 用 DTensor 让 AOTAutograd 不需要"识别 FSDP" —— 普通 trace 流程就能正确处理分片张量。

13.49 整章信息密度小结

读完 Ch 13 你应该能：

  • 理解：AOTAutograd 在编译栈的位置（连接 Dynamo 与 Inductor）
  • 追踪：joint trace → functionalize → partition 三段式工作流
  • 决策：何时用 dynamic shape、何时手动 activation_checkpoint vs 让 min-cut 自动决定
  • 诊断：TORCH_LOGS=aot_joint_graph,aot_graphs 看每段产物
  • 扩展：知道第三方 backend 怎么接进 AOTAutograd

AOTAutograd 是编译栈最复杂的中间环节，也是 PyTorch v2.0 演进路线上最大的工程投入之一。它把"动态图 PyTorch"与"静态优化编译"的鸿沟填上 —— 这是 PyTorch 2.x 编译栈能成立的关键一环。

13.50 设计启示

AOTAutograd 的核心思想可迁移：

第一：用图论算法（min-cut）做"显存 vs 计算"权衡：人手调 activation_checkpoint 是不可扩展的，让算法决定。这是编译器自动化的胜利

第二：先 functionalize 再优化：mutation 让分析变难，编译期把 mutation 抹掉、做完优化再恢复。这套"先简化、再优化、最后还原"是编译器中端经典模式

第三：joint graph 比独立 forward + backward 更优：在 joint 上做切分能找到全局最优，而不是 forward 与 backward 各自局部最优

第四：runtime wrapper 让编译产物兼容 eager：不需要破坏现有 autograd Engine，让 fw/bw 子图通过 autograd.Function 接入

下一章拆 Inductor —— AOTAutograd 把切好的子图送给 Inductor，Inductor 生成 Triton kernel。这是整个编译栈的"代码生成"环节。