快捷方式

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. 给定一些输入,它返回一个 FX 图,其中包含已执行的 PyTorch 函数 3. 如果检测到整数在调用之间发生变化,它也可以追踪整数 4. 它专门化任何其他不是张量或标量的值

当然,Dynamo 做了更多的事情,例如弄清楚何时需要重新追踪、重写函数的字节码、实现图中断……为了保持介绍简短,我们将在后续内容中逐步讨论所有这些。

PEP 523:向 CPython 添加帧评估 API

现在想象一下,我们被赋予了实现 Dynamo 的任务。我们甚至从哪里开始呢?对我们来说非常方便的是,PEP 523 随 Python 3.6 发布。这个 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 的工作是安装必要的脚手架,以便在调用函数时将字节码、args、全局变量等传递给 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 对象,并在内部保留 VariableTrackers 列表VariableTracker 的另一个示例是 ConstantVariable。ConstantVariable 包装了 Dynamo 认为恒定的对象。我们还有特殊子类用于需要特别注意的对象,例如 TensorVariable。所有这些内部类都在 torch/_dynamo/variables 文件夹中定义。

Python 对象在 VariableBuilder._wrap 中被包装到它们对应的 VariableTracker 类中。此函数只是一个非常长的 elif 链,它尝试将 Python 输入递归地模式匹配到适当类型的 VariableTracker 中。

调试技巧。当我们从 dynamo 获得意外结果时,有时是由构建器引起的。如果构建器的逻辑错误,有时 Dynamo 可能会将变量包装在不正确的 VariableTracker 类型中,这可能会在以后引起问题。查看错误中出现的 VariableTracker 类型以及在遇到 Dynamo 错误时抛出异常的 VariableTracker 方法非常有用。特别是,有时我们发现一个对象被跟踪为 UserDefinedObjectVariable(这是 Dynamo 的 catch-all 类),而它本应被跟踪为更具体的东西。在这些情况下,SourceBuilder.__call__ 逻辑通常是罪魁祸首。

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

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

这是原始程序的字节码和此时堆栈的状态。这对于查找对象未被追踪到正确的 VariableTracker 的位置非常有用。

好的,所以我们有一个追踪器的 IR,现在我们只需要重新实现 CPython 的栈式计算机。这是由 symbolic_convert.py 中的 InstructionTranslatorBase 实现的。

InstructionTranslatorBase 有大约 200 种方法,实现了几乎所有的 Python 字节码。例如,我们可以看到 BUILD_LIST 的实现

def BUILD_LIST(self, inst):
    items = self.popn(inst.argval)
    self.push(ListVariable(items, mutation_type=ValueMutationNew()))

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

生成输出图

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

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

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

使 Dynamo 健全:Guard

至此,我们有了一种完全忽略控制流来追踪程序的方法。为此,我们重新实现了所有的 CPython……如果这听起来有点杀鸡用牛刀,那是因为确实如此。torch.jit.trace 已经实现了这一点,而无需所有这些机制,那么问题出在哪里呢?

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

JAX 通过始终重新追踪并在重新追踪后缓存图来解决此问题。另一方面,Dynamo 使用 guard 来避免每次都重新追踪整个程序。

guard 是为了专门化一组示例输入的帧而做出的假设(输入上的布尔表达式)。只有当这些假设在新输入上成立时,重用图才是有效的。

例如,函数的任何常量输入(如字符串)都会安装一个 guard,声明该输入应为 str 类型,并且等于我们传递的字符串。运行

import torch

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

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

使用 TORCH_LOGS=guards 打印(以及其他 guard)

___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 来查看失败的 guard

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

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

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

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

xy 的源是 LocalSourcey[0] 的源是 GetItemSource,它在内部存储一个 LocalSource。另一方面,a 将没有源,因为它是一个仅存在于 fx 图中的中间变量。

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

import torch

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

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

生成以下 guard

___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])包装 LocalSourceL['l'])。

此时,有了源和 guard,我们就能够实现一个缓存系统,以避免重新编译,而无需每次都重新追踪。我们将在后续内容中更详细地讨论这个缓存系统。

细心的读者会注意到,这还没有解释为什么我们需要如此精细地控制 Python 解释器,以至于不得不重新实现它。我们展示的 guard 示例取决于输入对象,因此我们仍然可以在执行函数之前计算这些 guard。换句话说,我们可以基于 torch.jit.trace 实现此 guard 系统,并以更少的工作量获得相同的功能……输入符号形状。

符号形状

我们在介绍中讨论的另一点是,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,)

在第一个图中,形状被追踪为常量,但一旦它发生变化,它就会使用 SymInt 符号化地追踪它。一般来说,查看中间值的形状的更简单方法是使用 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]

我们看到在第一次调用时,guard 检查张量是否具有一些固定的尺寸和步幅。这些 guard 在第二次执行中失败,因此它会重新追踪。由于是 int guard 失败了,因此在第二次迭代中,它以符号方式追踪这个 int,并在更通用的内核上安装更通用的 guard。

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

0, 1 总是被专门化

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

做出此选择有几个原因。有两个特别重要 - 当且仅当张量的任何维度为零时,张量才为空 - 仅当步幅之一为一时,张量才能是连续的

此策略决策不适用于普通的 Python 整数;如果我们认为 Python 整数应该动态编译,我们默认不会专门化它们;相反,它是否被专门化取决于它的用法。

鸭子类型形状

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

符号整数上的 Guard

现在我们了解了符号形状是如何在高层次上实现的以及它们具有的属性。现在,为什么符号形状迫使我们走上获取 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 的 guard。就函数的输入而言,这是一个非平凡的 guard,但它是在程序执行过程中注册的。更重要的是,在我们看到以 SymNodeVariable 参数为条件的 if 语句之前,我们无法知道是否需要此 guard。这些条件对于 torch.jit.trace 是不可见的,并且需要对 python 代码进行深入分析。

调试提示 使用 TORCH_LOGS=dynamo 运行此代码会告诉我们 guard 添加的位置

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

在那里放置断点并查看回溯对于理解 guard 的来源非常有用。

使 Dynamo 完整:图中断

有了我们讨论的所有工具,我们有了一个追踪器,它可以追踪张量和整数上的 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 解释器运行编译后的代码。对于没有图中断的函数,调用该函数 2 次且参数相同的程序的整个追踪/执行过程如下所示

  1. 在第一次调用函数时

    1. Dynamo 将函数追踪到 FX 图中

      1. FX 图由编译器 (Inductor) 编译为高效的底层代码……但这又是另一天的故事了

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

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

  2. 在第二次调用函数时

    1. 它根据新参数检查第一次调用的 guard [这里]。由于它们与之前的参数相同,因此它们通过了

    2. 它要求 CPython 运行与这些 guard 关联的字节码 [这里]

这个过程本身看起来过于复杂。为什么生成新的字节码并要求 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 的综合开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源