目录
目录

编译的 Autograd：捕获更大的向后图torch.compile

创建时间： Oct 09， 2024 |上次更新时间：2024 年 10 月 23 日 |上次验证： Oct 09， 2024

作者： Simon Fan

您将学到什么

  • 编译后的 autograd 如何与torch.compile

  • 如何使用编译后的 autograd API

  • 如何使用TORCH_LOGS

先决条件

概述

Compiled Autograd 是 PyTorch 2.4 中引入的扩展 这允许捕获更大的向后图形。torch.compile

虽然 确实会捕获后向图，但它会部分捕获。AOTAutograd 组件提前捕获后向图，但存在某些限制：torch.compile

  • 向前的图形中断到向后的图形中断

  • 不捕获向后钩子

编译后的 Autograd 通过直接与 autograd 引擎集成来解决这些限制，从而允许 it 在运行时捕获完整的向后图。具有这两个特征的模型应该尝试 编译了 Autograd，并且可能会观察到更好的性能。

但是，Compiled Autograd 引入了自己的限制：

  • 在向后缓存查找开始时添加了运行时开销

  • 由于捕获较大，因此在 Dynamo 中更容易出现重新编译和图形中断

注意

Compiled Autograd 正在积极开发中，尚未与所有现有的 PyTorch 功能兼容。有关特定功能的最新状态，请参阅 已编译的 Autograd 登陆页面。

设置

在本教程中，我们将基于这个简单的神经网络模型进行示例。 它采用一个 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)

基本用法

在调用 API 之前，请确保设置为 ：torch.compiletorch._dynamo.config.compiled_autogradTrue

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)

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

  • Python 解释器调用 Dynamo，因为此调用用 .@torch.compile

  • Dynamo 截获 Python 字节码，模拟其执行并将作记录到图形中。

  • AOTDispatcher禁用钩子并调用 autograd 引擎来计算 和 的梯度，并将作记录到一个图表中。使用 AOTDispatcher 可重写 .model.linear.weightmodel.linear.biastorch.autograd.Functiontrain

  • Inductor 生成与 AOTDispatcher forward 和 backward 的优化实现相对应的函数。

  • Dynamo 将优化后的函数设置为接下来由 Python 解释器计算。

  • Python 解释器执行优化的函数，该函数执行 .loss = model(x).sum()

  • Python 解释器执行 ，调用 autograd 引擎，该引擎路由到编译的 Autograd 引擎，因为我们设置了 .loss.backward()torch._dynamo.config.compiled_autograd = True

  • 编译的 Autograd 计算 和 的梯度，并将作记录到一个图表中，包括它遇到的任何钩子。在此过程中，它将记录之前由 AOTDispatcher 重写的向后。然后，编译后的 Autograd 会生成一个新函数，该函数对应于 的完全跟踪实现，并在推理模式下执行该函数。model.linear.weightmodel.linear.biasloss.backward()torch.compile

  • 相同的步骤递归地应用于编译的 Autograd 图形，但这次 AOTDispatcher 不需要对图形进行分区。

检查编译的 autograd 日志

使用环境变量运行脚本：TORCH_LOGS

  • 要仅打印编译后的 autograd 图，请使用TORCH_LOGS="compiled_autograd" python example.py

  • 要打印具有更多张量元数据的图形并重新编译原因，以牺牲性能为代价，请使用TORCH_LOGS="compiled_autograd_verbose" python example.py

重新运行上面的代码段，编译后的 autograd 图形现在应该记录到 .某些图形节点的名称前缀为 、 这些对应于先前在 AOTAutograd 后向图 0 中提前编译的节点，例如，对应于 id=0 的 AOT 后向图。stderraot0_aot0_view_2view_2

在下图中，红色框封装了在没有 Compiled Autograd 的情况下捕获的 AOT 后向图。torch.compile

../_images/entire_verbose_log.png

注意

这是我们将调用 的图形，而不是优化的图形。编译后的 Autograd 实质上会生成一些未优化的 Python 代码来表示整个 C++ autograd 执行。torch.compile

使用不同的标志编译 forward 和 backward pass

你可以对这两个编译使用不同的编译器配置，例如，即使 forward 中有 graph break，backward 也可能是一个 fullgraph。

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)()

或者你可以使用上下文管理器，它将应用于其范围内的所有 autograd 调用。

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

编译的 Autograd 解决了 AOTAutograd 的某些限制

  1. 向前传递中的图形中断不再必然导致向后传递中的图形中断：

@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)

在第一种情况下，我们看到由于编译函数中的 2 个图形中断而产生了 3 个反向图形。 而在第二个使用 compiled autograd 的情况下，我们看到尽管有图中断，但还是追踪了一个完整的反向图。torch.compilefntorch.compile

注意

在跟踪 Compiled Autograd 捕获的向后弯钩时，Dynamo 仍有可能打断图形。

  1. 现在可以捕获向后钩子

@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()

图形中应该有一个节点，Dynamo 稍后会将该节点内联到以下内容中：call_hook

../_images/call_hook_node.png

Compiled Autograd 的常见重新编译原因

  1. 由于 loss 值的 autograd 结构发生了变化：

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")()

在上面的示例中，我们在每次迭代中调用不同的运算符，导致每次跟踪不同的 autograd 历史记录。您应该会看到一些重新编译的消息：Cache miss due to new autograd node。loss

../_images/recompile_due_to_node.png

  1. 由于张量改变形状：

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")()

在上面的示例中，更改形状，编译后的 autograd 将在第一次更改后标记为动态形状张量。您应该会看到 recompiles 消息：Cache miss due to changed shapes。xx

../_images/recompile_due_to_dynamic.png

结论

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

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源