快捷方式

Dynamo 深入解读

TorchDynamo(或简称为 Dynamo)是 torch.compile 中的跟踪器,它通常是导致那些疯狂的回溯的罪魁祸首。但是,我们不能盲目地责怪 Dynamo 出现这些错误。为了向用户提供灵活的功能,Dynamo 被赋予了理解任何 Python 程序的艰巨任务。特别是,Dynamo 必须在内部实现 Python 编程语言的大部分功能!

在这篇文章中,我们将从头开始介绍 Dynamo 的内部设计。我们将讨论它提供的功能以及它的实现方式。在阅读完这篇文章后,您将更好地理解当您 torch.compiled 一个 PyTorch 程序并导致编译错误,或者编译成功但速度提升不符合预期时,到底哪里出了问题。

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_code python example.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 会向自定义解释器提供以下信息:- 函数的字节码 - 函数参数的值(即局部变量)及其名称 - 全局变量的值及其名称 - 内置函数,例如 absprint

您可以在 这里查看所有字段。 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。ConstantVariable 包装了 Dynamo 认为是常量的所有对象。我们还有针对需要特殊关注的对象的特殊子类,例如 TensorVariable。所有这些内部类都定义在 torch/_dynamo/variables 文件夹中。

Python 对象在 VariableBuilder._wrap 中被包装到相应的 VariableTracker 类中。这个函数只是一长串 elif,试图递归地将 Python 输入与适当类型的 VariableTracker 进行模式匹配。

调试技巧。当我们从 dynamo 中得到意外的结果时,有时是由于构建器造成的。如果构建器的逻辑错误,有时 Dynamo 可能会将变量包装到错误的 VariableTracker 类型中,这可能会导致以后出现问题。查看错误中出现的 VariableTracker 类型和在遇到 Dynamo 错误时抛出异常的 VariableTracker 方法非常有用。特别是,有时我们会发现某个对象被跟踪为 UserDefinedObjectVariable(这是 Dynamo 的万能类),而它应该被跟踪为更具体的类型。在这些情况下,通常是 SourceBuilder.__call__ 逻辑出了问题。

调试技巧。在使用 TORCH_LOGS=dynamo 运行程序时,打印出来的工件之一是以下形式的行

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

这是原始程序的字节码,以及此时堆栈的状态。这对于找出对象没有被跟踪到正确的 VariableTracker 的位置非常有用。

好的,所以我们有了跟踪器的 IR,现在我们只需要重新实现 CPython 的栈式机器。这是由 InstructorTranslatorBasesymbolic_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。这意味着我们从堆栈顶端弹出前 3 个元素,并将由这三个元素组成的新的列表对象压入堆栈顶部。

生成输出图

有了符号执行 Python 代码的方法,我们就可以提取在给定一些输入的程序的符号执行期间发生的 PyTorch 操作。这在 Dynamo 中通过 OutputGraph 对象实现。 OutputGraph 对象 绑定到 `InstructionTranslator 对象,它跟踪创建 FX 图所需的所有数据,这些数据将由 Dynamo 返回。

FX 图的所有输入和中间元素都是 fx.Node。在 Dynamo 中,fx.Node 被包装在 fx.Proxy 中。 fx.Proxy 用于构建 FX 图。特别是,它们将对它们执行的每个 PyTorch 操作记录到图中。您可以通过调用 create_proxy 创建一个要添加到图中的新操作。然后,我们可以通过函数 wrap_fx_proxy 将其添加到图中。

图存储对张量的操作……以及对符号整数的操作。我们将在后面讨论符号整数,但首先我们将讨论 Dynamo 如何解决一个相当重要的正确性问题。

使 Dynamo 健壮:守卫

到目前为止,我们已经找到了完全忽略控制流来跟踪程序的方法。为此,我们重新实现了整个 CPython……如果这听起来有点过分,那是因为它确实如此。 torch.jit.trace 已经实现了这一点,而不需要所有这些机制,那么问题出在哪里呢?

torch.jit.trace 的问题在于,正如其文档中所警告的那样,它只有在跟踪的程序不是数据相关的时才有效。换句话说,只有在程序本身是线性的情况下它才能正常工作。这意味着编写程序时不能使用 if-else 语句、for-while 循环、异常。更重要的是,我们使用的所有库都不能使用任何控制流!总之,在像 Python 这样动态的语言中不使用控制流实际上是一个巨大的限制。

JAX 通过在每次重新跟踪后重新跟踪并缓存图来解决这个问题。另一方面,Dynamo 使用守卫来避免每次都重新跟踪整个程序。

守卫是一种假设(对输入的布尔表达式),用于专门针对一组示例输入对框架进行专门化。只有当这些假设在新的输入上成立时,重用图才有效。

例如,对函数的任何常量输入,例如字符串,都会安装一个守卫,说明该输入应该为 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'

函数的输入在构建器中包装 以及 在程序执行期间,守卫会累积起来。 我们将在下一节中展示更多守卫示例,但首先让我们讨论来源。

**来源**跟踪如何从进入当前帧时存在的原始局部或全局变量中重建变量。 特别是,它跟踪原始局部和全局对象以及它们包含的任何对象。 在

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

xy 的来源为 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 图中记录对它的所有操作。 4 我们已经在介绍中介绍符号整数跟踪时看到了此类。

现在让我们讨论定义 Dynamo 中符号形状跟踪的三个属性,以及如何实现它们。

默认情况下为静态

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

使用 TORCH_LOGS=graph_code 运行此程序,我们看到这两个调用被跟踪为

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

在第一个图中,形状被跟踪为常量,但一旦它发生变化,它就会使用 SymInts 符号地跟踪它。 通常,查看中间值形状的一种更简单方法是使用 TORCH_LOGS=graph_sizes 运行程序

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)

我们可以看到,两个张量参数的第一维是动态的,因为它由 s0 变量表示。

我们可以通过运行 TORCH_LOGS=guards 来找到 Dynamo 如何实现这一点。

# 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]

我们看到在第一次调用中,守卫检查张量是否具有一些固定的大小和跨距。 这些守卫在第二次执行中失败,因此它重新跟踪。 由于这是一个失败的 int 守卫,因此在第二次迭代中,它会符号地跟踪此 int,并在该更通用的内核上安装更通用的守卫。

**编译性能提示**。 如果你知道一个维度的大小会变化,你可以在调用 torch.compile 之前调用 torch._dynamo.mark_dynamic 来将其标记为动态。 这将避免使用静态形状进行第一次编译。 还有其他有用的实用程序函数,如 maybe_mark_dynamicmark_static。 你也可以通过调用 torch.compile(dynamic=True) 来跟踪所有整数和形状。 这主要用于调试目的。

0、1 始终专业化

无论我们是否将一个维度标记为动态,如果我们传递一个输入,其中该维度为 0 或 1,Dynamo 将将其跟踪为非动态,并且它将为其生成一个特定的图。 这就是为什么在上面的示例中我们找到形式为 2 <= L['a'].size()[0] 的守卫的原因。

这种选择有几个原因。 两个特别重要的原因是 - 当且仅当张量的一个维度为零时,该张量为空 - 张量只能在其中一个跨距为一时才是连续的。

此策略决定不适用于普通的 Python ints;如果我们认为 Python int 应该动态编译,我们不会默认情况下对其进行专业化;相反,它是否进行专业化取决于它的使用方式。

鸭子形状

Dynamo 执行我们所谓的“鸭子形状”。 如果两个动态整数在跟踪时具有相同的值,我们将假设它们相等并对其进行守卫。 事实上,这意味着,与其在上面的示例中具有两个符号 s0s1,不如将它们统一为 s0,并使用守卫 L['b'].size()[0] == L['a'].size()[0]。 这使得能够在编译器中执行融合,同时能够生成足够通用的内核。

符号 ints 上的守卫

现在我们了解了符号形状如何在高级别上实现,以及它们具有的属性。 现在,为什么符号形状迫使我们走上了控制 CPython 解释器的困难路线? 请考虑以下示例。

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

此代码有一个形式为 2*L['a'].size()[0] >= 16 的守卫。 这是一个在函数输入方面非平凡的守卫,但它是在程序执行过程中间注册的。 更重要的是,在我们看到以 SymNodeVariable 参数为条件的 if 语句之前,我们无法知道需要此守卫。 这些条件对 torch.jit.trace 是不可见的,需要对 Python 代码进行深入分析。

**调试提示** 使用 TORCH_LOGS=dynamo 运行此代码会告诉我们此守卫是在哪里添加的。

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 绑定。 Dynamo 不可能跟踪通过 C++ 实现的 Python 对象。 当跟踪器遇到它不理解的操作时,它该怎么办?

机器学习跟踪器通常处理此问题的做法是通知用户它遇到了什么操作,然后放弃跟踪。 这在 PyTorch 的情况下会造成一个真正的可用性问题,因为 PyTorch 的用户习惯了它带来的灵活性。 作为现实世界中的例子,doctr_det_predictor 模型使用 NumPy 和 cv2 库来 后处理模型的结果

另一个需要访问 CPython 的地方就在这里。Dynamo 可以让 CPython 运行有问题的代码,而不是报错!为此,Dynamo 在跟踪时会生成两个图:一个包含问题代码之前的所有操作,另一个包含问题代码之后的所有操作。 6 然后,在运行时,它会委托 CPython 执行第一个图,然后执行有问题的代码,最后执行第二个图。这种停止跟踪并生成多个图的过程称为 **图断点**。

一个小小的坦白:我在介绍和第一部分中撒了谎。Dynamo 并不是生成一个图,而是 **多个图**!实际上,在第二个图之后重新开始跟踪可以被视为开始跟踪一个新的函数。图断点后的新图将拥有它自己的保护、新的局部变量集等等。

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

  1. 在第一次调用该函数时

    1. Dynamo 将该函数跟踪到一个 FX 图

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

    2. 它重写了该函数的字节码,使其只调用编译后的函数

    3. 它将新的字节码交给 CPython 并要求其运行 [这里]

  2. 在第二次调用该函数时

    1. 它检查第一次调用中的保护是否与新的参数匹配 [这里]。由于它们与之前的参数相同,所以它们通过了

    2. 它要求 CPython 运行与这些保护相关的字节码 [这里]

这个过程本身看起来过于复杂。为什么生成新的字节码并让 CPython 运行它,而不是简单地创建一个指向编译后的函数的 C++ 绑定并执行它?好吧,这种模式使我们能够实现图断点!图断点生成的字节码具有以下结构

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

  • 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 对象的形式,因此上面的调试技巧有时也有助于找出是什么导致了图断点。

结论

Dynamo 是一款复杂的软件。一旦你注册了要实现一个 CPython 解释器,你就知道自己要开始一段旅程了。也就是说,我们希望这篇文章能帮助你解开一些谜团。

Dynamo(主要)是用 Python 实现的。我们留下了许多指向我们讨论过的代码部分的链接。我们希望阅读这些代码部分并搜索调用它们的地方,或者在它们上设置断点并查看调用栈,有助于理解其余的代码库。

当然,学习软件工作原理的最佳方式是扩展它。在这种情况下,最好的方法是查看 GitHub 上的 Dynamo 开放问题。它们中的许多只需要对代码进行很小的修改,一旦你找到了需要修改的地方。

脚注

1

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

2

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

3

在 CPython 行话中,所有这些对象的集合被称为 **框架**。框架

4

还有 SymBoolSymFloat 类。在撰写本文时,后者并没有被广泛使用。

5

有趣的是,它确实理解 NumPy 代码!请查看 这篇博文文档。现在,这之所以有可能,是因为我们使用 PyTorch 重新实现了 NumPy。不过,祝你在 PyTorch 中实现 Django 好运……

6

假设只有一段有问题的代码。如果有多段,Dynamo 可以根据需要将代码分成多个图。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源