快捷方式

动态形状

代码: symbolic_shapes.py

另请参阅: 动态形状手册

动机

深度学习编译器通常仅适用于静态形状,也就是说,它们生成的编译程序仅适用于输入形状的单个特定配置,并且如果任何输入形状发生更改,则必须重新编译。这种假设对于当今运行的大多数常见深度学习模型都非常有效,但在某些情况下,它是不够的

  • 某些维度(例如批次大小或序列长度)可能会有所不同。例如,执行自适应批处理的推理服务将根据其批处理窗口中收到的请求数量,以不同的批次大小执行推理请求。我们可能还希望仅将可变大小的序列填充到批次内的最大序列长度,这可能因批次而异。

  • 某些模型表现出数据相关的输出形状,也就是说,其输出和中间值的大小可能取决于实际输入数据,而实际输入数据可能因运行而异。例如,检测模型可能会首先生成可变数量的潜在边界框,然后再运行更昂贵的图像识别模型,以识别主体是否在边界框中。边界框的数量是数据相关的。

  • 当处理稀疏表示(例如稀疏张量、锯齿状张量和图神经网络)时,会发生一种特别重要的数据相关形状的情况。在所有这些情况下,要处理的数据量取决于问题的稀疏结构,而稀疏结构通常会以数据相关的方式变化。

在支持动态形状时,我们选择不支持动态秩程序,例如,输入张量维度发生变化的程序,因为这种模式在现实世界的深度学习程序中很少发生,并且它可以避免对形状的符号列表进行归纳推理的需要。

简略的公共 API

PyTorch 2.1 中的默认动态行为是

  • PT2 默认假定一切都是静态的

  • 如果我们因为大小更改而重新编译,我们将尝试将该大小重新编译为动态大小(已更改的大小将来可能还会更改)。这种泛化可能会失败(例如,因为用户代码在有问题的尺寸上进行了条件分支,或者 PT2 中缺少动态形状支持)。如果您想了解为什么 PT2 过度专业化了某些代码,请使用 TORCH_LOGS=dynamic 运行,并查找 “eval” 条目,其中说明了何时添加保护以及原因。

  • 如果您提前知道某些内容将是动态的,则可以使用 torch._dynamo.mark_dynamic(tensor, dim) 跳过第一次重新编译。如果您提前知道此维度可以采用的 minmax 值,则可以指定 torch._dynamo.mark_dynamic(tensor, dim, min=min, max=max)

  • 如果您说 torch.compile(dynamic=False),我们将关闭重新编译时的自动动态形状,并始终为每个不同的大小重新编译。相反,如果您说 torch.compile(dynamic=True),我们将尝试使一切尽可能动态化。这主要适用于小型运算符;如果您在大模型上尝试,它将 (1) 可能会使 PT2 崩溃,并且 (2) 无缘无故地运行缓慢。

保护模型

在考虑如何向 TorchDynamo 和 TorchInductor 添加动态形状支持时,我们做出了一个重要的设计决策:为了重用以 Python/C++ 编写的针对 PyTorch API 的分解和其他预先存在的代码,我们必须能够追踪动态形状。与可以捕获条件分支的两个分支的完全符号系统不同,我们总是选择一个分支并在假设我们仅在将来对该分支做出相同选择时才使用此跟踪的情况下专门化我们的跟踪。为此,我们为每个符号大小维护一个“提示”,说明其在编译时的具体值(由于 TorchDynamo 是一个即时编译器,因此它始终知道实际的输入大小)。当我们对张量执行条件操作时,我们只需查阅提示即可找出要采取哪个分支。

这大大简化了我们生成的符号形状公式,但也意味着我们有一个更为复杂的系统来管理保护。例如,考虑以下程序

def f(x, y):
    z = torch.cat([x, y])
    if z.size(0) > 2:
        return z.mul(2)
    else:
        return z.add(2)

我们将使用 TorchInductor 编译的最终 IR 将是 torch.cat([x, y]).add(2)torch.cat([x, y]).mul(2)(条件被展平),但是为了确定我们处于哪个分支,我们需要知道中间值 z 的大小。由于 TorchDynamo 必须预先知道编译后的跟踪是否有效(我们不支持像某些 JIT 编译器那样的退出),因此我们必须能够将 z.size(0) 简化为输入的表达式,x.size(0) + y.size(0)。这是通过为 PyTorch 中的所有运算符编写元函数来完成的,这些运算符可以将大小信息传播到张量的输出,而无需实际对节点执行计算。

总体架构

符号形状工作流程

  1. 当我们开始在 Dynamo 中编译帧时,我们分配一个 ShapeEnv(附加到 FakeTensorMode),它跟踪符号形状状态。

  2. 我们在输入时为张量分配符号大小(什么是静态的或动态的是策略决策,带有一些旋钮)。

  3. 我们通过运算符传播符号大小,同时维护 (1) FX IR,以便我们可以忠实地导出符号计算,以及 (2) 表示大小变量的 Sympy 表达式,以便我们可以推理它们。

  4. 当我们对符号大小进行条件判断时,无论是在 Dynamo 跟踪中还是在 Inductor 优化中,我们都会根据条件添加保护。这些可以从 Python 和 C++ 中推导出来。

  5. 这些保护可以进一步简化符号变量。例如,如果您断言 s0 == 4,我们现在可以将所有出现的 s0 替换为 4

  6. 当我们完成跟踪和优化后,我们会将所有这些保护与编译后的代码一起安装;只有当所有保护评估为真时,编译后的代码才是可重用的。

重要文件

  • C++ SymInt API: c10/core/SymInt.h, SymFloat.h, SymBool.h

  • Python SymInt API: torch/__init__.py (查找 SymInt/SymFloat/SymBool)

  • C++ 管道: c10/core/SymNodeImpl.h, torch/csrc/utils/python_symnode.h, torch/csrc/jit/python/init.cpp

  • Python 基础设施: torch/fx/experimental/symbolic_shapes.py

  • 其他重要文件: torch/_subclasses/fake_tensor.py, torch/_meta_registrations.py, decomps, PrimTorch refs

简略的内部 API

理解 Python 类层次结构

  • SymInt/SymFloat/SymBool: 这些是用户可见的类,它们模拟其 int/float/bool 对等物。如果您添加两个 SymInt,我们将为您提供一个新的 SymInt,它以符号方式跟踪发生的整数加法。

  • SymNode: 这是内部结构(可通过例如 symint.node 访问),它保存实际的符号跟踪信息。 SymNode 是类型擦除的;这使得表示混合类型操作更加方便。请注意,从技术上讲,您不必从 SymInt 调用到 Python SymNode;例如,XLA 的 C++ SymNodeImpl 将取代 SymNode 的位置。

  • ShapeEnv: 每个编译上下文状态,用于跟踪到目前为止我们累积的所有自由符号和保护。每个 SymNode 记录其 ShapeEnv(但反之亦然;SymNode 仅在它们参与保护时才使用)。

C++ 非常相似

  • c10::SymInt/SymFloat/SymBool: 用户可见的类,它们模拟 int/float/bool。

  • c10::SymNode/SymNodeImpl: 类似于 SymNode

  • C++ 中没有 ShapeEnv;为了便于调试,整个符号推理装置都在 Python 中。

当您编写可以使用 make_fx 跟踪的代码时,它必须能够处理流经它的 SymInt/SymFloat/SymBool。动态形状手册 提供了一些关于如何执行此操作的指导。

DimDynamic 策略

符号推理

  • 值范围

  • Sympy 用法说明

  • 约束

  • DimDynamic/Constraint

无后备 SymInt

为了解决控制流,我们检查符号整数的提示(即实际值)以确定要转到哪个分支。但是,在某些情况下,我们可能没有提示:所谓的无后备符号整数出现在大小变量从数据相关操作(如 .nonzero().item())中出现时。对这些符号整数执行控制流是非法的,因此我们必须在这些操作上进行图形中断。

朴素地实现,这太严格了:如果您尝试对无后备符号整数执行任何操作,大多数 PyTorch 程序都会立即失败。以下是使之真正起作用的最重要的增强功能

  • 在张量创建时,PyTorch 会预先计算有关张量的大量数据;例如,如果您使用 empty_strided 创建张量,我们将急切地对步幅进行排序,并确定张量是否是非重叠的和密集的。排序会产生大量的保护。但是,更常见的是使用更高级别的 API(如 empty)直接生成张量,这保证生成非重叠且密集的张量。我们修改了 PyTorch 以避免不必要地重新计算这些属性。

  • 即使需要进行重要的计算,有时也根本不会查询属性。使这些预计算属性延迟加载使我们能够避免在无后备符号整数上进行保护,除非实际需要它。

  • 整数张量中的数据通常已知为非负数。但是,我们提供了一个 API constrain_range,用户可以使用该 API 指定大小以已知限制为上限和下限。

在 PT2 的未来版本(PT2.1 之后)中,我们将扩展我们的推理系统,以根据用法推断无后备符号整数是类似大小的。例如,如果您将 .item() 调用的结果传递给工厂函数(如 torch.empty),我们将自动推断结果是一个大小(因为如果不是,它将失败。)此假设将在运行时得到验证,如果未满足,则会引发错误。

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源