第 1 课：PyTorch 执行栈、Dispatcher 与 Autograd

本课目标与完成标准

本课产出：从 Python op 追到 dispatcher/kernel/autograd node，并实现 DAG 反向中的梯度累加。

通过要求：唯一代码填空题通过给定检查；Q1～Q3 都给出因果链；能指出一个正确性边界和一个性能/运营取舍；总分至少 8/10。

与既有六条主线的边界

cuda/ 和 triton/ 解释 kernel 内部；本课解释框架如何选择 kernel、记录动态图并在反向汇总多条梯度路径。

前置：Python、PyTorch、train 第 1～5 课、CUDA 基础。本课只补 AI Infra 缺口，不重复已经完成的 kernel 或并行算法推导。

核心概念与运行机制

1. 它是什么，解决什么问题

当我们在 Python 中执行 torch.matmul(a, b) 或 torch.add(x, y) 时，PyTorch 从 Python 顶层调用到底层硬件发射对应的 CUDA 内核，中间跨越了极其复杂的底层运行时栈（Runtime Stack）。这一调用栈的核心由两大基础设施支撑：c10::Dispatcher 分发器 与 Autograd 自动微分引擎。

1. c10::Dispatcher 动态分发中心：
   • 解决的核心问题是多维度解耦与异构后端的零开销多路复用。
   • 一个张量算子在实际执行时，面临着多维度的正交状态：运算在 CPU 还是 CUDA 设备上？数据是 FP32 还是 INT8 量化类型？是否需要记录反向图（Autograd）？是否处于批量矢量化执行模式（vmap / Batched）？是否处于图编译追踪模式（Tracer / FakeTensor）？
   • 如图 1-1 所示，c10::Dispatcher 摒弃了

Python API 接口层

torch.matmul(a, b) / torch.add(x, y)

c10::Dispatcher 多键动态路由

Dispatch Key: Autograd | CPU | CUDA | Tracer | Batched

底层 C++ / CUDA 实现

Autograd::AddBackward0 (构造反向图)

at::native::cuda_kernel() (GPU 执行)

图 1-1：PyTorch Dispatcher 调度键分发与执行栈拓扑结构

在算子内部写死庞大 if-else 分支的高耦合设计，采用高效的位图调度键集（DispatchKeySet）与函数指针查找表，将分发时延压缩在数十纳秒级别，实现了异构设备、量化、自动求导与编译器的无缝插拔。

2. Autograd 动态计算图引擎：
   • 解决的核心问题是深度学习动态网络的前向自动图构建与反向高效链式拓扑求导。
   • 无论前向包含多少动态分支（Python if / while 循环），Autograd 均能以即时执行（Eager Execution）的方式按需生成包含反向梯度算子节点（Node / Function）的有向无环图（DAG），并在反向调用 loss.backward() 时精准完成梯度的多路聚合与叶子节点回传。

2. 它如何工作

PyTorch 运行时的执行流涵盖了从分发器多键解包到反向就绪队列遍历的完整数据流与控制流：

1. Dispatcher 的多键拆包与降级转发（Dispatch Pipeline）：
   • 如图 1-1 所示，当 Python 算子被调用时，系统根据输入张量的属性计算出当前的综合键集 DispatchKeySet（例如同时包含 DispatchKey::AutogradCUDA 与 DispatchKey::CUDA）；
   • Dispatcher 按照严格预设的优先级栈遍历分发键。若张量标记为 requires_grad=True，分发器首先命中最顶层的 AutogradCUDA 处理函数；
   • 在该函数内部，Autograd 引擎创建对应的反向节点（如 AddBackward0），并将前向计算所需的张量（Saved Tensors）通过钩子打包保存，建立起计算图的边连接关系；
   • 随后，系统调用 redispatch 剥离当前处理过的 Autograd 键，分发器自动降级路由到次高优先级的 DispatchKey::CUDA，调用 ATen 底层 C++ 算子 at::native::add_cuda 真正发射 GPU 计算内核。
2. Autograd 反向动态拓扑遍历与 ReadyQueue（Backward Workflow）：
   • 如图 1-2 所示，当用户调用 loss.backward() 时，

标量损失 loss.backward()

待处理反向节点工作队列 (ReadyQueue)

弹出当前就绪 Node (如 MulBackward0)

执行反向算子计算输入叶子张量梯度

更新或累加 .grad 属性 (AccumulateGrad)

检查并递减前驱边依赖计数

将前驱 Node 入队

等待其他路径汇合

就绪未完成

图 1-2：Autograd 动态计算图反向拓扑遍历与梯度累加工作流

引擎在后台启动一个专用的反向执行引擎（Engine::execute），并初始化工作队列 ReadyQueue； - 损失函数的根节点首先被推入 ReadyQueue； - 调度线程持续从 ReadyQueue 中弹出已就绪的 Node（例如 MulBackward0），解包前向保存的张量，调用其 apply() 方法计算输入张量的局部梯度； - 若梯度回传的目标张量是用户定义的模型参数（Leaf Variable），触发 AccumulateGrad 算子，原子更新该参数的 .grad 属性； - 若梯度流向前驱中间节点，引擎递减该前驱节点的入度依赖计数。只有当前驱节点的所有下游依赖分支全部完成计算且梯度在节点处完成相加汇合时，该前驱节点才正式就绪并被压入 ReadyQueue，直至整张图遍历完毕。

3. 正确性条件与常见误区

• 版本计数器与原地修改检测（In-Place Modification & Version Counter）：
  • 核心设计：每个 Tensor 底层由 c10::TensorImpl 管理，内部维护一个原子递增的整型变量 version_counter_；
  • 安全契约：若张量被正向算子保存以备反向求导（Saved for Backward），引擎记录其当时的基准版本号（saved_version）。若在正向计算之后、反向求导之前，用户对该张量执行了原地修改（In-Place Operation，如 x += 1、x.relu_() 或切片赋值 x[0] = 0），其 version_counter_ 将发生递增；
  • 当反向执行到对应算子解包张量时，系统触发严格断言： $$\text{tensor.\_version} == \text{saved\_version}$$ 若不相等，立即抛出 RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation，严厉杜绝使用被篡改的数据导致静默数学错误。
• 叶子节点梯度累加契约（Gradient Accumulation Invariant）：
  • 在 Autograd 设计中，参数的 .grad 字段在每次反向求导时是执行原位相加（In-Place Addition）而非直接覆盖。
  • 这一契约是支撑大模型分布式显存受限时实现“微批次梯度累加（Gradient Accumulation over Micro-Batches）”的物理根基；但在标准训练循环中，必须在每个 Optimizer Step 之后显式调用 optimizer.zero_grad(set_to_none=True)，否则跨 Step 的梯度污染将导致模型发散。
• no_grad() vs. inference_mode() vs. detach()：
  • torch.no_grad()：在上下文内屏蔽 Autograd 调度键，正向不构建图，但保留版本计数器更新机制；
  • torch.inference_mode()：比 no_grad() 更激进，完全禁用版本计数器与 View 追踪机制，极大压低开销，但产出的张量禁止离开上下文参与反向；
  • tensor.detach()：创建与原张量共享底层物理 Storage 的新浅拷贝视图，但将 grad_fn 置为 None，斩断与原计算图的拓扑联系。

4. 性能与工程取舍

• 前向激活保存 vs. 激活重计算（Activation Storing vs. Checkpointing）：
  • 全量保存激活：反向无需重复执行正向算子，计算时间最短；但显存占用随序列长度与 Transformer 层数呈线性甚至二次方剧烈膨胀，迅速导致 OOM；
  • 激活重计算（Activation Checkpointing / Rematerialization）：在正向阶段仅保存层间边界张量，丢弃昂贵的中间注意力激活矩阵；在反向阶段按需重新执行前向算子。全量重计算相当于多做一遍前向，即前向 FLOPs 增加约 100%，按“反向约为前向 2 倍”估算约占训练总 FLOPs 的 +33%。以此为代价，释放高达 70% 的激活显存，是大模型训练不可或缺的工程杠杆。
• Eager 模式调度开销 vs. torch.compile 图编译：
  • 在 Eager 模式下，即便底层 CUDA 内核仅耗时 1 微秒，Python 绑定、参数解析与 Dispatcher 多层 C++ 分发的开销通常也有约 5~10 微秒，导致 GPU 计算流频繁遭遇 CPU 调度气泡；
  • torch.compile（基于 TorchDynamo、AOTAutograd 与 Inductor）在图级别捕获计算子图，将多个细粒度算子融合为单个 Triton 内核，完全绕过 Eager Dispatcher 的逐算子发射循环，显著提升算力占空比。

具体演示

设标量损失由两条分支汇合而成：$y = f(x) + g(x) = x^2 + x^3$。

1. 前向传播：记录计算图，输入 $x$ 同时连接至 $x^2$ 分支与 $x^3$ 分支；

2. 反向求导链式汇聚：

   • 沿第一条分支传回局部梯度：$\frac{\partial (x^2)}{\partial x} = 2x$；
   • 沿第二条分支传回局部梯度：$\frac{\partial (x^3)}{\partial x} = 3x^2$；

3. 多路梯度在叶子节点 $x$ 处执行拓扑累加： $$\frac{\partial y}{\partial x} = \left.\frac{\partial (x^2)}{\partial x}\right|_{x} + \left.\frac{\partial (x^3)}{\partial x}\right|_{x} = 2x + 3x^2$$ 设 $x = 2.0$： $$\left.\frac{\partial y}{\partial x}\right|_{x=2} = 2(2) + 3(2^2) = 4 + 12 = 16.0$$ 若调度系统存在覆盖赋值 Bug（只保留最后一条边），梯度将被错误计算为 $12.0$，导致参数更新偏离最优下降方向。

实践任务：唯一代码填空题

补齐多条反向边对同一 tensor 的梯度累加。

只能修改 TODO/______ 位置，不得删除断言或放宽通过条件。

def accumulate_edge_grads(edge_grads):
    """edge_grads: [(tensor_name, local_grad), ...]"""
    totals = {}
    for name, grad in edge_grads:
        # TODO：同一 tensor 可能出现多次，不能覆盖已有贡献。
        totals[name] = ______
    return totals

assert accumulate_edge_grads([("x", 4.0), ("x", 12.0), ("w", 3.0)]) == {"x": 16.0, "w": 3.0}
assert accumulate_edge_grads([]) == {}

检查方法

运行本单元格，所有断言必须通过；再补一个边界输入并解释预期。

Q1

为什么调用两次 loss.backward() 默认会让 parameter.grad 相加，而不是覆盖？

你的答案：

Q2

对 forward 保存的 tensor 做原地修改，为什么 autograd 需要版本计数？

你的答案：

Q3

自定义 op 有 CUDA kernel 但 torch.compile/autograd 不认识，集成至少需要哪些层？

你的答案：

评分规则

• 代码 4 分：正常输入 2 分、边界输入 1 分、解释 1 分；
• Q1～Q3 各 2 分；
• 一票否决：混淆测量与推测、忽略租户/请求隔离、用平均值掩盖尾延迟、删除失败路径。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先独立完成。核对后请改变一个规模或故障条件重新推演。

def accumulate_edge_grads(edge_grads):
    """edge_grads: [(tensor_name, local_grad), ...]"""
    totals = {}
    for name, grad in edge_grads:
        totals[name] = totals.get(name, 0.0) + grad
    return totals

assert accumulate_edge_grads([("x", 4.0), ("x", 12.0), ("w", 3.0)]) == {"x": 16.0, "w": 3.0}
assert accumulate_edge_grads([]) == {}

Q1 参考答案

叶子参数由 AccumulateGrad 节点执行 grad.add_，所以多次 backward 会累加。这支撑梯度累积与多 loss 训练，代价是每个 step 后要 zero_grad(set_to_none=True)。

Q2 参考答案

Autograd 必须维护原子版本计数器（version_counter_）以杜绝原地修改带来的隐蔽致命数学错误：

1. 求导公式对前向状态的刚性依赖（Mathematical Dependency）：深度学习中绝大多数非线性算子的导数表达式严格依赖前向传播时的局部状态数值。例如：$\text{ReLU}(x)$ 的反向掩码取决于前向时的 $x > 0$；$\text{Softmax}(x)$ 的导数直接依赖于前向生成的输出激活 $y$。Autograd 通过 save_for_backward 机制将这些张量的引用保存在反向节点中；

2. 原地修改的破坏本质（In-Place Mutation Hazard）：若用户或框架内部执行了原地操作（In-Place Operation，如 x.add_(1)、x.relu_() 或切片赋值 x[0] = 5），张量底层 Storage 中的物理数据被直接改写。若没有版本校验，反向传播执行时将直接读取被篡改后的脏数据计算梯度。更危险的是，程序不会抛出任何系统崩溃，而是在输出完全错误的虚假梯度下继续训练，导致模型精度严重衰减或无法收敛；

3. 版本计数器的闭环防御机制：每个 Tensor 底层的 c10::TensorImpl 均持有一个原子自增整型 version_counter_。任何原地修改都会触发版本号自增。反向节点解包张量时强制比对快照版本：若 current_version != saved_version，立即 Fail-Fast 抛出 RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation，确保数学正确性的物理确定性。

Q3 参考答案

将一个具备 CUDA Kernel 的自定义 C++ 算子无缝接入 PyTorch 并获得 Autograd 及 torch.compile（TorchDynamo / Inductor）原生支持，必须完整构建以下四层接口契约：

1. C10 Dispatcher 算子声明与后端注册（Schema & CUDA Kernel Registration）：
   • 使用 TORCH_LIBRARY(custom_ops, m) 定义算子的函数签名与强类型 Schema（如 custom_ops::my_gemm(Tensor a, Tensor b) -> Tensor）；
   • 使用 TORCH_LIBRARY_IMPL(custom_ops, CUDA, m) 将底层 C++/CUDA Kernel 绑定到 DispatchKey::CUDA，使 Eager 模式能够完成设备分发与张量解包；
2. Autograd 自动求导分发绑定（Autograd Function / DispatchKey::Autograd）：
   • 若在 Python 侧集成，继承 torch.autograd.Function 实现静态 forward 与 backward 方法，手动管理 ctx.save_for_backward 与输入梯度的推导；
   • 若在 C++ 层面原生集成，使用 TORCH_LIBRARY_IMPL(custom_ops, AutogradCUDA, m) 注册专用的反向分发函数，动态构建反向 Node 并建立计算图拓扑边；
3. Meta / FakeTensor 内核注册（Shape & Dtype Propagation）：
   • 使用 TORCH_LIBRARY_IMPL(custom_ops, Meta, m) 注册 Meta Kernel。该函数仅依据输入张量的 Shape、Stride 与 Dtype 快速推导出输出张量的几何尺寸与内存布局，不分配任何 GPU 显存也不发射 CUDA Kernel。这是 torch.compile 在 AOTAutograd 追踪与符号推演（Symbolic Shape Analysis）阶段的核心先决条件；
4. Inductor 编译器 Lowering 与算子解构（Decomposition & Lowering to Triton）：
   • 为算子定义分解规则（register_decomposition），指示 AOTAutograd 如何将其解构为已有的基础 ATen 原语；或者在 PyTorch Inductor 后端注册专用的 Triton 代码生成模板（register_lowering），使编译器能够将该自定义算子与前后的 LayerNorm 或激活函数进行跨算子全局内核融合（Kernel Fusion）。

参考资料

• PyTorch autograd mechanics
• PyTorch C++ Dispatcher internal
• PyTorch Custom Operator Registration Guide
• PyTorch AOTAutograd & Inductor Compiler

API 与平台能力会演进；部署前应按目标版本重新核对。

← 上一课
第 8 课：Transform Dialec… 下一课 →
第 2 课：torch.compile