Dynamo 深度解析#

  • admin
  • 2025-12-08 11:21:28

使 Dynamo 完整:图中断(Graph Breaks)#

有了我们讨论过的所有工具,我们就拥有了一个追踪器,它可以追踪张量和整数上的 PyTorch 操作,并拥有一个知道何时可以重用先前追踪的图以及何时需要重新追踪的缓存系统。所有这一切都是在执行任意 Python 代码!

这里只有一个小问题。陈述“执行任意 Python 代码”可能有点太笼统了。Dynamo 实现了一大部分 Python,但它是否实现了协程或异步等更复杂的部分?它是否实现了整个 Python 标准库?NumPy 也有一个 Python API。 torch.compile 是否也理解 NumPy?以及 Django?5

Python 的生态系统非常庞大,其中很大一部分是用 C++ 或 Rust 等更高效的语言编写的,它们只是暴露了 Python 绑定。Dynamo 无法追踪用 C++ 实现的 Python 对象。当追踪器遇到它不理解的操作时,它能做什么?

机器学习追踪器处理这个问题的常规方法是告知用户它们卡住的操作,然后完全放弃追踪。在 PyTorch 的情况下,这会带来真正的可用性问题,因为它的用户习惯了它提供的灵活性。例如,doctr_det_predictor 模型使用 NumPy 和 cv2 库来后处理模型的输出。

这里是 CPython 的另一个有趣之处。Dynamo 不是抛出错误,而是可以让 CPython 运行那个有问题的代码!为了做到这一点,Dynamo 在追踪时生成一个图,包含有问题代码之前的所有操作,以及一个包含有问题代码之后的所有操作的图。6 然后,在运行时,它会将执行第一个图、有问题的代码以及第二个图的任务委托给 CPython。这种停止追踪并生成多个图的过程称为图中断。

一个小小的坦白:我在整个介绍和前几节都说了谎。Dynamo 不只生成一个图,而是生成多个图!对于所有实际目的,在第二个图之后开始重新追踪可以被认为是在开始追踪一个新函数。图中断后的新图将有自己的 Guard,自己的一组局部变量,等等。

为了讨论如何实现图中断,我们首先需要回顾一下 Dynamo 如何与 CPython 交互。使用 PEP 523,CPython 允许用户使用自己的帧评估机制。我们还没有讨论的是,CPython 还公开了自己的帧评估供他人使用。Dynamo 利用这一点让快速的 CPython 解释器运行编译后的代码。对于没有图中断的函数,程序第一次和第二次调用函数(参数相同)的整个追踪/执行过程如下:

在第一次调用函数时

Dynamo 将函数追踪成一个 FX 图

FX 图由编译器(Inductor)编译成高效的低级代码…但这又是另一个故事了

它重写函数的字节码,使其只需调用编译后的函数

它将这个新的字节码提供给 CPython,并要求它运行它在此处

在第二次调用函数时

它将第一次调用的 Guard 与新参数进行比较在此处。由于参数与之前相同,因此通过

它要求 CPython 运行与这些 Guard 相关联的字节码在此处

这个过程本身看起来过于复杂。为什么还要生成新的字节码并要求 CPython 运行它,而不是直接创建一个 C++ 绑定到编译后的函数并执行它?嗯,这个模式允许我们实现图中断!由图中断生成的字节码具有以下结构:

执行第一个图的字节码

离开堆栈的字节码,就像 CPython 执行第一个图时的状态一样。它还会重放当时可见的对局部或全局变量的任何修改

导致 Dynamo 图中断的字节码

执行第二个图的字节码

让我们看一个简单的例子

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(原始函数)和一个名为 resume_in_fn 的函数。第二个函数是由 Dynamo 创建的,用于实现从图中断开始的程序执行。这通常被称为续延函数 (continuation function)。此续延函数只需使用正确的参数调用第二个编译后的函数。初始函数的代码被重写,实现了我们之前描述的策略:

L0-4. 调用编译后的函数(a + 2)。

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

L8-18. 离开堆栈,使其处于图中断时的状态

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

L22-32. 调用编译后的续延函数(a + b)

Dynamo 中堆栈的代码生成委托给了 VariableTracker 子类。Dynamo 中的每个 VariableTracker 对象都有一个 reconstruct 方法,该方法生成必要的字节码以在堆栈上创建它所代表的 Python 对象。

调试技巧。图中断会影响性能,因此最好避免它们。使用 TORCH_LOGS=graph_breaks 运行程序是查找程序触发了多少图中断的一个好方法。它返回的信息是关于 VariableTracker 对象,所以上面的调试技巧有时也有助于弄清楚是什么原因导致了该图中断。