编译自动求导：捕获更大的反向传播图以进行 torch.compile()

编译自动求导：捕获更大的反向传播图以进行 `torch.compile()`#

编译自动微分（Compiled Autograd）是 PyTorch 2.4 中引入的 torch.compile() 扩展，它能够捕获更大的反向传播图。

尽管 torch.compile() 确实捕获了反向传播图，但它只是部分捕获。AOTAutograd 组件提前捕获反向传播图，但存在一些限制：

前向传播中的图断裂会导致反向传播中的图断裂
反向传播钩子未被捕获

编译自动微分通过直接与自动微分引擎集成，解决了这些限制，使其能够在运行时捕获完整的反向传播图。具有这两种特性的模型应该尝试使用编译自动微分，并可能观察到更好的性能。

然而，编译自动微分也引入了自身的限制：

在反向传播开始时增加了缓存查找的运行时开销
由于捕获的图更大，在 dynamo 中更容易导致重新编译和图断裂

备注

编译自动求导功能正在积极开发中，尚未与所有现有的 PyTorch 功能兼容。要了解特定功能的最新状态，请参阅编译自动求导功能登陆页面。

在本教程中，我们将基于这个简单的神经网络模型进行示例演示。该模型接收 10 维的输入向量，通过线性层进行处理，并输出另一个 10 维的向量。

import torch

class Model(torch.nn.Module):
   def __init__(self):
      super().__init__()
      self.linear = torch.nn.Linear(10, 10)

   def forward(self, x):
      return self.linear(x)

基本用法#

在调用 torch.compile() API 之前，请确保将 torch._dynamo.config.compiled_autograd 设置为 True：

model = Model()
x = torch.randn(10)

torch._dynamo.config.compiled_autograd = True
@torch.compile
def train(model, x):
   loss = model(x).sum()
   loss.backward()

train(model, x)

/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()
/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()

在上述代码中，我们创建了 Model 类的一个实例，并通过 torch.randn(10) 生成了一个随机的 10 维张量 x。我们定义了训练循环函数 train，并使用 @torch.compile 装饰它以优化其执行。当调用 train(model, x) 时：

Python 解释器调用 Dynamo，因为此调用被 @torch.compile 装饰。
Dynamo 拦截 Python 字节码，模拟其执行并记录操作到图中。
AOTDispatcher 禁用钩子并调用 autograd 引擎来计算 model.linear.weight 和 model.linear.bias 的梯度，并将操作记录到图中。使用 torch.autograd.Function，AOTDispatcher 重写了 train 的前向和反向实现。
Inductor 生成一个与 AOTDispatcher 前向和反向的优化实现相对应的函数。
Dynamo 设置优化后的函数，以便 Python 解释器接下来对其进行求值。
Python 解释器执行优化后的函数，该函数执行 loss = model(x).sum()。
Python 解释器执行 loss.backward()，调用 autograd 引擎，由于我们设置了 torch._dynamo.config.compiled_autograd = True，因此路由到编译自动求导引擎。
编译自动求导引擎计算 model.linear.weight 和 model.linear.bias 的梯度，并将操作记录到图中，包括它遇到的任何钩子。在此过程中，它将记录之前由 AOTDispatcher 重写的反向操作。编译自动求导引擎随后生成一个新函数，该函数对应于 loss.backward() 的完全追踪实现，并在推理模式下使用 torch.compile 执行它。
相同的步骤递归地应用于编译自动求导图，但这次 AOTDispatcher 不需要对图进行分区。

检查编译自动求导日志#

通过设置 TORCH_LOGS 环境变量运行脚本：

仅打印编译自动求导图，使用 TORCH_LOGS="compiled_autograd" python example.py
打印包含更多张量元数据和重新编译原因的图，但会牺牲性能，使用 TORCH_LOGS="compiled_autograd_verbose" python example.py

重新运行上述代码片段，编译自动求导图现在应记录到 stderr。某些图节点将带有前缀 aot0_，这些节点对应于之前在 AOTAutograd 反向图 0 中提前编译的节点，例如，aot0_view_2 对应于 id=0 的 AOT 反向图中 view_2。

在下图中，红色框包含由 torch.compile 捕获的 AOT 反向图，而未使用编译自动求导。

备注

这是我们将调用 torch.compile 的图，不是优化后的图。编译自动求导本质上生成一些未优化的 Python 代码来表示整个 C++ 自动求导执行过程。

使用不同标志编译前向和反向传播#

你可以为两次编译使用不同的编译器配置，例如，即使前向传播中存在图断裂，反向传播也可能是全图。

def train(model, x):
    model = torch.compile(model)
    loss = model(x).sum()
    torch._dynamo.config.compiled_autograd = True
    torch.compile(lambda: loss.backward(), fullgraph=True)()

或者，你可以使用上下文管理器，它将应用于其作用域内的所有自动求导调用。

def train(model, x):
   model = torch.compile(model)
   loss = model(x).sum()
   with torch._dynamo.compiled_autograd.enable(torch.compile(fullgraph=True)):
      loss.backward()

编译自动求导解决了 AOTAutograd 的某些限制#

前向传播中的图断裂不再必然导致反向传播中的图断裂：

@torch.compile(backend="aot_eager")
def fn(x):
   # 1st graph
   temp = x + 10
   torch._dynamo.graph_break()
   # 2nd graph
   temp = temp + 10
   torch._dynamo.graph_break()
   # 3rd graph
   return temp.sum()

x = torch.randn(10, 10, requires_grad=True)
torch._dynamo.utils.counters.clear()
loss = fn(x)

# 1. base torch.compile
loss.backward(retain_graph=True)
assert(torch._dynamo.utils.counters["stats"]["unique_graphs"] == 3)
torch._dynamo.utils.counters.clear()

# 2. torch.compile with compiled autograd
with torch._dynamo.compiled_autograd.enable(torch.compile(backend="aot_eager")):
   loss.backward()

# single graph for the backward
assert(torch._dynamo.utils.counters["stats"]["unique_graphs"] == 1)

在第一个 torch.compile 案例中，由于编译函数 fn 中的 2 个图断裂，生成了 3 个反向图。而在第二个使用编译自动求导的 torch.compile 案例中，尽管存在图断裂，仍然追踪到了一个完整的反向图。

备注

当追踪由编译自动求导捕获的反向钩子时，Dynamo 仍然可能发生图断裂。

现在可以捕获反向钩子

@torch.compile(backend="aot_eager")
def fn(x):
   return x.sum()

x = torch.randn(10, 10, requires_grad=True)
x.register_hook(lambda grad: grad+10)
loss = fn(x)

with torch._dynamo.compiled_autograd.enable(torch.compile(backend="aot_eager")):
   loss.backward()

计算图中应该有 call_hook 节点，Dynamo 稍后会将其内联到以下内容中：

编译自动求导的常见重新编译原因#

由于损失值的自动求导结构发生变化：

torch._dynamo.config.compiled_autograd = True
x = torch.randn(10, requires_grad=True)
for op in [torch.add, torch.sub, torch.mul, torch.div]:
   loss = op(x, x).sum()
   torch.compile(lambda: loss.backward(), backend="eager")()

在上面的例子中，我们在每次迭代时调用不同的操作符，导致 loss 每次都跟踪不同的自动求导历史。你应该会看到一些重新编译的消息：由于新的自动求导节点导致的缓存未命中。

由于张量形状的变化：

torch._dynamo.config.compiled_autograd = True
for i in [10, 100, 10]:
   x = torch.randn(i, i, requires_grad=True)
   loss = x.sum()
   torch.compile(lambda: loss.backward(), backend="eager")()

在上面的例子中，x 的形状发生了变化，编译后的自动求导会在第一次变化后将 x 标记为动态形状的张量。你应该会看到重新编译的消息：由于形状变化导致的缓存未命中。

结论#

在本教程中，我们探讨了 torch.compile 与编译后的自动求导的高级生态系统，介绍了编译后的自动求导的基础知识，以及一些常见的重新编译原因。请继续关注 dev-discuss 的深入探讨。