Dynamo 深度解析

Dynamo 入门介绍

在深入了解所有实现细节之前，让我们先来讨论一下Dynamo的作用。

Dynamo 是一个追踪器。它会执行给定的函数和输入，并将一系列线性的、不含控制流的指令记录到图中。例如，考虑以下程序：

import torch

@torch.compile
def mse(x, y):
    z = (x - y) ** 2
    return z.sum()

x = torch.randn(200)
y = torch.randn(200)
mse(x, y)

如果我们把这个程序保存到 example.py 文件中，并运行

TORCH_LOGS=graph_codepythonexample.py

我们看到了Dynamo的追踪输出

def forward(l_x_: torch.Tensor, l_y_: torch.Tensor):
    # File: example.py:5, code: z = (x - y) ** 2
    sub = l_x_ - l_y_
    z = sub ** 2
    # File: example.py:6, code: return z.sum()
    sum_1 = z.sum()
    return (sum_1,)

我们将这称为给定输入下的函数图（或跟踪）。这是通过FX 图来表示的。我们也可以把 FX 图简单地看作是一个存储函数调用列表的容器。

首先，我们应该注意到，该图展示了一个线性的 PyTorch 操作序列。1 Dynamo 记录所有的 PyTorch 操作并按顺序存储它们。例如，它将 z = (x - y) ** 2 分解为两个操作：sub = l_x_ - l_y_ 和 z = sub ** 2。

当我们说跟踪是线性的时，意味着没有分支和控制流的存在。为了更好地理解这一点，请考虑以下情况：

import torch

@torch.compile
def fn(x, n):
    y = x ** 2
    if n >= 0:
        return (n + 1) * y
    else:
        return y / n

x = torch.randn(200)
fn(x, 2)

当使用TORCH_LOGS=graph_code执行时，它会返回

def forward(l_x_: torch.Tensor):
    # File: example.py:5, code: y = x ** 2
    y = l_x_ ** 2
    # File: example.py:7, code: return (n + 1) * y
    mul = 3 * y
    return (mul,)

我们看到，Dynamo 完全从追踪中移除了 if 语句，仅记录了实际执行的操作及其输入。

因此，应该清楚函数的跟踪依赖于输入。特别是这意味着，在编写@torch.compile时并不会生成跟踪，而是在执行带有实际参数的函数fn(x, 2)时才会生成。

这里值得注意的是，Dynamo 去掉了函数的第二个参数，并将其视为常量。它将操作结果 n + 1 记录在图中。这是 Dynamo 的一个特性：除了整数之外，Dynamo 将所有非张量值都视为常量。现在让我们来看看整数有什么特别之处。

Dynamo 最后一个定义性属性是它能够处理动态形状。具体来说，Dynamo 能够追踪形状和整数的变化，而不是将它们视为常量。这使得可以避免重新编译，并部署适用于任意大小的通用模型以供生产使用。动态形状主要出现在批量大小中，在训练时我们可能使用固定批量大小，但在推理阶段会遇到任意大小的批量。此外，处理文本或音频时也会遇到可变序列长度。

我们可以通过多次运行上面的示例来观察这一现象

import torch

@torch.compile
def fn(x, n):
    y = x ** 2
    if n >= 0:
        return (n + 1) * y
    else:
        return y / n

x = torch.randn(200)
fn(x, 2)
fn(x, 3)
fn(x, -2)

在这种情况下，设置 TORCH_LOGS=graph_code 将生成两个额外的图表

# Graph for n==2 omitted

def forward(self, l_x_: torch.Tensor, l_n_: torch.SymInt):
    # File: a.py:5, code: y = x ** 2
    y = l_x_ ** 2

    # File: a.py:7, code: return (n + 1) * y
    add = l_n_ + 1
    mul = add * y
    return (mul,)

def forward(self, l_x_: torch.Tensor, l_n_: torch.SymInt):
    # File: a.py:5, code: y = x ** 2
    y = l_x_ ** 2

    # File: a.py:9, code: return y / n
    truediv = y / l_n_
    return (truediv,)

Dynamo 检测到一个整数在第一次调用后发生了值的变化，并开始对其进行跟踪。我们看到这些图是通用的，并通过类型为 SymInt 的对象对变量 n 进行符号追踪。

如果我们在这几次调用后调用 fn(x, 4)，Dynamo 不会重新编译，而是重用之前已记录的图。

总结如下：1. Dynamo 是一个 Python 跟踪器；2. 给定一些输入，它会返回执行的 PyTorch 函数的 FX 图形；3. 如果检测到整数在调用之间发生变化，它可以跟踪这些整数；4. 它可以专门处理除张量或标量以外的任何其他值。

当然，Dynamo 还执行了许多其他任务，例如确定何时需要重新跟踪、重写函数的字节码以及实现图中断等。为了使介绍简明扼要，我们将在这之后的部分中逐步讨论所有这些内容。

PEP 523：在CPython中添加帧评估API

假设我们现在需要实现 Dynamo，我们应该从哪里开始呢？幸运的是，随着 Python 3.6 的发布，PEP 523 出现了。这个 PEP 设计目的 是为了允许第三方为 Python 创建 JIT 编译器。让我们来看看具体如何操作。

关于CPython的说明: CPython 内部实现为一个栈式机器。Python 程序会被编译成字节码，然后由解释器执行这些字节码。要了解更多关于字节码的信息，请参阅标准库中的 dis 模块。另见开发者文档，了解 CPython 解释器的介绍。我们将假设读者熟悉栈式机器的概念。

PEP 523 提供了一个 API，允许用户为每个函数添加自定义解释器。然后，CPython 使用这个自定义解释器而不是默认的来执行该函数。为了能够执行函数，在进入时，CPython 向自定义解释器提供以下内容：- 函数的字节码 - 函数参数（即局部变量）及其名称的值 - 全局变量及其名称的值 - 内置函数如 abs 或 print

你可以在这里查看所有字段。2

总之，CPython 为用户解释器提供了执行函数所需的全部信息。3

通过这个API，我们可以实现一个追踪器，即通过实现一个解释器来运行代码并记录所有在此执行过程中发生的PyTorch操作。这正是Dynamo所做的事情。

Dynamo 使用这个 CPython API 来解析所有对象，并将它们打包进一个 Python 结构中。完成这些操作后，它会从 C 语言返回到 Python。除了这段与 CPython 通信的代码外，Dynamo 其余部分完全用 Python 实现。

应该清楚的是，装饰器 @torch.compile 的职责是安装必要的框架结构，在函数调用时将字节码、参数和全局变量传递给 Dynamo。同样地，@torch.compile 并不实际执行编译。

用 Python 实现 CPython

因此，我们回到了 Python 世界。我们有了一个函数的字节码，并且具备了执行它的所有上下文信息。特别是，我们在_convert_frame_assert 函数处着陆。这是装饰器 torch.compile 返回的函数！我们通过_dynamo.optimize 函数到达此函数。装饰器 torch.compile 只是为 _dynamo.optimize 提供的一个友好 API。

在实现 Python 解释器之前，我们希望定义一个中间表示（IR）。具体来说，我们将把所有局部和全局变量包装到我们自己的内部类中。这样可以更好地跟踪这些对象，并将可以以相同方式处理的对象组合在一起，以便 Dynamo 能够更有效地管理它们。

VariableTracker 是内部类结构的父类，它表示 Dynamo 理解的不同对象。例如，ListVariable 表示一个 list 对象，并在内部维护一个VariableTracker 列表。另一个 VariableTracker 的例子是ConstantVariable，它包装了所有 Dynamo 认为是常量的对象。我们还为需要特别关注的对象定义了特殊的子类，例如TensorVariable。所有这些内部类都在torch/_dynamo/variables 文件夹中定义。

Python 对象会被封装到对应的 VariableTracker 类中，具体实现位于 VariableBuilder._wrap。此函数通过一系列的 elif 语句递归地将 Python 输入匹配到适当的 VariableTracker 类型。

调试技巧. 当从 Dynamo 获取到意外结果时，这可能是由 builder 引起的。如果 builder 的逻辑有误，Dynamo 可能会将变量包装在不正确的 VariableTracker 类型中，从而导致后续问题。当你遇到 Dynamo 错误时，检查错误中的 VariableTracker 类型以及抛出异常的方法是非常有用的。特别是有时我们发现一个对象被错误地追踪为通用的 UserDefinedObjectVariable 类型，而它应该被更具体地追踪。在这种情况下，问题通常在于 SourceBuilder.__call__ 逻辑。

调试技巧。当你使用TORCH_LOGS=dynamo运行程序时，会打印出类似以下格式的日志条目

TRACE LOAD_GLOBAL y [TorchInGraphFunctionVariable(<built-in method any>), TensorVariable()]

这是原始程序的字节码及其当时的栈状态。这非常有助于找到对象未能正确追踪到相应VariableTracker的位置。

好的，我们现在有了一个用于追踪器的中间表示（IR），接下来只需要重新实现 CPython 的栈机器。这在 InstructorTranslatorBase 中实现，位于 symbolic_convert.py。

InstructionTranslatorBase 大约包含 200 个方法，实现了几乎所有 Python 字节码的功能。例如，我们可以查看 BUILD_LIST 的具体实现。

def BUILD_LIST(self, inst):
    items = self.popn(inst.argval)
    self.push(ListVariable(items, mutable_local=MutableLocal()))

这是类似l = [2, 3, 4]构造生成的字节码。在这种情况下，由于有三个元素，生成的字节码是BUILD_LIST 3。这意味着我们将栈顶的三个元素弹出，并将由这三个元素组成的新的列表对象推送到栈顶。

生成输出图形

有了对Python代码进行符号执行的方法，我们现在可以提取在给定输入的情况下，在程序的符号执行过程中发生的PyTorch操作。这通过Dynamo中的OutputGraph对象实现。OutputGraph 对象绑定到一个 InstructionTranslator 对象，并跟踪创建FX图所需的所有数据，该图将由Dynamo返回。

FX 图中的所有输入和中间元素都是 fx.Nodes。在 Dynamo 中，fx.Nodes 被包裹在 fx.Proxys 中。fx.Proxys 用于构建 FX 图，并记录在其上执行的每个 PyTorch 操作。你可以通过调用create_proxy 创建一个新的操作，并通过函数wrap_fx_proxy 将其添加到图中。

图存储了对张量和符号整数的操作。我们稍后会详细介绍符号整数，但首先我们将讨论Dynamo如何解决一个重要正确的性问题。

让Dynamo发出声音：守卫

到目前为止，我们已经找到了一种完全忽略控制流来追踪程序的方法。为此，我们重新实现了整个 CPython……如果这听起来有点过于复杂，那确实如此。torch.jit.trace 已经实现这一功能而不需要所有这些机制，那么问题来了？

torch.jit.trace 存在的问题，正如其文档警告的那样，仅当跟踪程序不依赖于数据时才有效。换句话说，如果程序本身是线性的，则可以正常工作。这意味着编写程序时不使用 if-else 语句、for 或 while 循环以及异常处理。此外，我们使用的任何库也不能包含控制流。总之，在像 Python 这样动态的语言中避免使用控制流实际上是一个巨大的限制。

JAX 通过在每次重新跟踪后缓存图形解决了这个问题。而 Dynamo 使用 guards 来避免每次都重新跟踪整个程序。

一个守卫是为了针对一组示例输入对框架进行特殊化处理而作出的假设（即输入上的一个布尔表达式）。只有在新输入满足这些假设的情况下，重用该图才是有效的。

例如，对于函数中的任何常量输入（如字符串），会设置一个守卫，声明该输入应该是类型str，并且与传递的字符串相等。

import torch

@torch.compile
def fn(a, b):
    return a * len(b)

fn(torch.arange(10), "Hello")

使用 TORCH_LOGS=guards 会打印各种保护措施（其中包括其他保护措施）

___check_type_id(L['b'], 94334122025024)
L['b'] == 'Hello'

这表明局部变量 b 应该具有特定类型（在这种情况下是 str，由常量 9433... 表示），并且其值应该是 'Hello'。如果我们再次执行函数并传递不同的参数

import torch

@torch.compile
def fn(a, b):
    return a * len(b)

fn(torch.arange(10), "Hello")
fn(torch.arange(10), "Hi")

我们可以通过运行TORCH_LOGS=recompiles来查看失败的保护措施

Recompiling function fn in script.py:3
triggered by the following guard failure(s):
     - L['b'] == 'Hello'

当函数输入被构建器包装时 和 程序执行期间，guards 会被累积。在下一节中我们将展示更多 guards 的示例，在此之前先来讨论一下 sources。

一个源记录了如何从进入当前作用域时存在的原始局部和全局变量中重新构建某个变量的方法。具体来说，它会追踪这些原始的局部和全局对象及其所包含的所有其他对象。

def foo(x: Tensor, y: List[Tensor]):
    a = x * y[0]
    return a * x

x 和 y 的来源是LocalSource，而 y[0] 的来源是GetItemSource，它内部包含了一个 LocalSource。另一方面，a 不会有来源信息，因为它是一个仅在 fx 图形中使用的中间变量。

所有这些都在torch/_dynamo/source.py 中定义。我们可以在以下示例中看到 GetItemSource 生成的保护措施：

import torch

@torch.compile
def fn(x, l):
    return x * len(l[0])

fn(torch.randn(8), ["Hi", "Hello"])

生成以下防护代码

___check_type_id(L['l'], 94439025877664)
len(L['l']) == 2
___check_type_id(L['l'][0], 94439025840192)
L['l'][0] == 'Hi'
___check_type_id(L['l'][1], 94439025840192)
L['l'][1] == 'Hello'

在这里，我们可以看到GetItemSource生成的代码（[0] 和 [1]）包装了一个LocalSource（L['l']）。

此时，借助源码和保护机制，我们可以实现一个缓存系统，从而避免每次都进行重新编译。我们将在后续部分更详细地介绍这个缓存系统。

细心的读者可能会注意到，这还没有解释为什么我们需要对Python解释器进行如此精细的控制，以至于必须重新实现它。我们展示的守卫示例依赖于输入对象，因此在执行函数之前仍然可以计算这些守卫。换句话说，我们可以在 torch.jit.trace 之上实现这个守卫系统，并且只需较少的努力就能获得相同的功能……引入符号形状。

符号形状

我们在介绍中提到的一个要点是 Dynamo 能够追踪整数。为了实现这一点，我们使用了一个名为 torch.SymInt 的符号类，它像一个普通的 int 类型一样工作，但会记录所有对其执行的操作，并在输出的 FX 图中显示。当我们在介绍符号整数追踪时已经见过这个类。

我们现在来讨论定义 Dynamo 中符号形状追踪的三个属性，以及如何实现这些属性。

import torch

@torch.compile
def fn(a, b):
    return a.shape[0] * a * b

fn(torch.randn(4, 3), torch.randn(4, 3))
fn(torch.randn(8, 3), torch.randn(8, 3))

def forward(self, l_a_: torch.Tensor, l_b_: torch.Tensor):
    mul = 4 * l_a_
    mul_1 = mul * l_b_
    return (mul_1,)

def forward(self, s0: torch.SymInt, l_a_: torch.Tensor, l_b_: torch.Tensor):
    size = l_a_.size()
    getitem = size[0]
    mul = getitem * l_a_
    mul_1 = mul * l_b_
    return (mul_1,)

TRACED GRAPH TENSOR SIZES
===== __compiled_fn_1 =====
l_a_: (s0, 3)
l_a_ (concrete): (8, 3)
l_b_: (s0, 3)
l_b_ (concrete): (8, 3)
mul: (s0, 3)
mul (concrete): (8, 3)
mul_1: (s0, 3)
mul_1 (concrete): (8, 3)

# Guards first call
check_tensor(L['a'], torch.float32, device=None, requires_grad=False, size=[4, 3], stride=[3, 1])
check_tensor(L['b'], torch.float32, device=None, requires_grad=False, size=[4, 3], stride=[3, 1])

# Guards second call
check_tensor(L['a'], torch.float32, device=None, requires_grad=False, size=[None, 3], stride=[3, 1])
check_tensor(L['b'], torch.float32, device=None, requires_grad=False, size=[None, 3], stride=[3, 1])

L['b'].size()[0] == L['a'].size()[0]
2 <= L['a'].size()[0]

import torch

@torch.compile(dynamic=True)
def fn(a):
    if a.shape[0] * 2 < 16:
        return a
    else:
        return a + 1

fn(torch.randn(8))

eval 2*s0 >= 16 [guard added] at script.py:5 in fn (_dynamo/variables/tensor.py:812 in evaluate_expr)

完善Dynamo：图中断

借助我们讨论过的所有工具，我们有一个能够追踪 PyTorch 张量和整数操作的跟踪器，并且它具有一个缓存系统，可以判断何时重用之前已追踪的图，以及何时需要重新进行追踪。这一切都是通过执行任意 Python 代码来实现的。

这里有一个小问题。“执行任意Python代码”这个说法可能有些过于笼统。Dynamo实现了一部分Python功能，但是否实现了更复杂的特性，比如协程或异步操作呢？它是否完整地实现了整个Python标准库？NumPy也有一个Python API。torch.compile 是否也支持NumPy？还有Django呢？5

Python 的生态系统非常庞大，其中很大一部分是用其他更高效的语言（如 C++ 或 Rust）编写的，并且只是提供了 Python 绑定。通过在 C++ 中实现的 Python 对象进行 Dynamo 跟踪是不可能的。当跟踪器遇到无法理解的操作时，它可以做什么？

通常，机器学习跟踪器遇到此类问题时会告知用户并停止追踪。在 PyTorch 中，这将是一个真正的可用性问题，因为其用户习惯于它的灵活性。例如，doctr_det_predictor 模型使用 NumPy 和 cv2 库来后处理模型的结果。

这里有一个地方，能够访问CPython是很有趣的。当遇到有问题的代码时，Dynamo不会出错，而是会将执行委托给CPython来运行这段代码。为了实现这一点，Dynamo在跟踪阶段生成两个图：一个包含问题代码之前的所有操作，另一个包含问题代码之后的所有操作。6 在运行时，它首先让CPython执行第一个图中的操作，然后执行有问题的代码，最后执行第二个图中的操作。这个过程中断跟踪并生成多个图的过程被称为图中断。

一个小 confession：我在介绍和前几节中一直在撒谎。Dynamo 不生成一个图，而是多个图！从实用角度来看，在第二个图之后开始重新追踪可以视为开始追踪一个新的函数。在图断开后的新的图会有它自己的守卫、新的局部变量集等。

要讨论如何实现图中断，我们需要先回顾 Dynamo 是如何与 CPython 交互的。根据 PEP 523，CPython 允许用户使用自己的帧评估机制。此外，CPython 还公开了自己的帧评估机制供他人使用。Dynamo 利用这一点让快速的 CPython 解释器运行编译后的代码。对于没有图中断的函数，在调用该函数两次且参数相同的情况下，整个跟踪和执行过程如下所示：

1. 在首次调用函数时

   1. Dynamo 将函数转化为 FX 图形

      1. FX 图由编译器（Inductor）编译成高效的低级代码……但这些细节我们另述

   2. 它修改了函数的字节码，使其直接调用编译后的版本

   3. 它给 CPython 提供了新的字节码，并要求其运行 [此处]

2. 在第二次函数调用时

   1. 它会将第一次调用时的guards与新的参数进行对比 [这里]。由于这些参数与之前的一致，因此通过了检查。

   2. 它要求 CPython 运行与这些守卫相关的字节码 [此处]

这个过程本身看起来过于复杂。为什么不去直接为编译后的函数创建C++绑定并执行它，而是要生成新的字节码并让CPython运行呢？这种模式的好处是可以实现图中断！由图中断生成的字节码具有以下结构：

1. 用于执行第一个图形的字节码

2. 生成的字节码会将堆栈保持在与 CPython 执行第一个图后的相同状态，并且还会重播此时可见的所有局部和全局变量的修改。

3. 导致Dynamo图形故障的字节码

4. 用于执行第二个图形的字节码

让我们通过一个简单的例子来看看这一点

import torch

@torch.compile
def fn(a):
    b = a + 2
    print("Hi")
    return b + a

fn(torch.randn(4))

使用 TORCH_LOGS=bytecode 运行会显示初始字节码和修改后的字节码

MODIFIED BYTECODE fn script.py line 3
 0 LOAD_GLOBAL              1 (__compiled_fn_0)
 2 LOAD_FAST                0 (a)
 4 CALL_FUNCTION            1
 6 STORE_FAST               3 (graph_out_0)
 8 LOAD_GLOBAL              0 (print)
10 LOAD_CONST               2 ('Hi')
12 LOAD_FAST                3 (graph_out_0)
14 LOAD_CONST               3 (0)
16 BINARY_SUBSCR
18 STORE_FAST               1 (b)

20 CALL_FUNCTION            1
22 LOAD_GLOBAL              2 (__resume_at_14_1)
24 ROT_TWO
26 LOAD_FAST                0 (a)
28 LOAD_FAST                1 (b)
30 CALL_FUNCTION            3
32 RETURN_VALUE

MODIFIED BYTECODE resume_in_fn script.py line 6
 0 LOAD_GLOBAL              1 (__compiled_fn_2)
 2 LOAD_FAST                2 (b)
 4 LOAD_FAST                1 (a)
 6 CALL_FUNCTION            2
 8 UNPACK_SEQUENCE          1
10 RETURN_VALUE

我们可以看到修改后的字节码被拆分为两个函数：一个是原始函数 fn，另一个是由 Dynamo 生成的名为 resume_in_fn 的续函数。这个续函数用于实现从图中断处开始执行程序的功能，它只是使用正确的参数调用第二个编译后的函数。初始函数的代码被重写，以实现我们之前描述的策略。

• L0-4. 调用编译后的函数（a + 2）。

• L6. 将其结果存储在名为 graph_out_0 的局部变量中。 graph_out_0 是一个元组。

• L8-18. 将堆栈状态保持在图中断点时的状态

• L20. 执行导致图表中断的代码

• L22-32. 调用编译后的续函数（a + b）

Dynamo 中堆栈的代码生成由 VariableTracker 子类负责。每个 VariableTracker 对象都有一个 reconstruct 方法，用于生成创建其表示的 Python 对象所需的字节码。

调试技巧：图中断会影响性能，应尽量避免。使用TORCH_LOGS=graph_breaks运行程序可以统计出我们的程序遇到了多少次图中断。返回的信息以VariableTracker对象的形式呈现，因此上面提到的调试技巧有时也有助于找出导致图中断的原因。

结论

Dynamo 是一个复杂的软件。一旦你决定实现一个 CPython 解释器，就知道这将是一个挑战性的旅程。尽管如此，我们希望这篇文章能够帮助大家更好地理解它。

Dynamo 主要是用 Python 编写的。我们提供了许多与讨论内容相关的代码片段的链接。希望通过阅读这些代码片段，并通过查找调用它们的位置，或在这些位置设置断点并检查调用堆栈，能够帮助理解整个代码库。

当然，了解软件如何工作的最佳方式是通过扩展它。在这种情况下，最好的方法是查看GitHub上开放的Dynamo问题。一旦找到需要修改的地方，许多问题只需要对代码进行很小的改动。

脚注

1

在文献中，这被称为有向无环图(DAG)。

2

所有的绑定代码都位于 torch/csrc/dynamo/eval_frame.c。

3

在 CPython 中，所有这些对象的集合被称为帧。

4

还有 SymBool 和 SymFloat 类。目前，SymFloat 类用得并不多。

5

有趣的是，它确实支持 NumPy 代码！请参阅这篇博客文章《编译 NumPy 代码》和相关文档《Torch 编译常见问题解答》。不过，这只是因为我们使用 PyTorch 重新实现了 NumPy。要在 PyTorch 中实现 Django 就另当别论了……

6

假设只有一个有问题的代码段。如果有多个问题，Dynamo 可以将代码拆分成交互的多个图。