目录

常见问题

在高层次上,TorchDynamo堆栈由使用dynamo从Python代码捕获的图和后端编译器组成。在这个示例中,后端编译器由使用AOTAutograd进行反向图追踪和使用TorchInductor进行图降低组成。当然还有许多其他编译器在这里可用,但对于本文档,我们将以inductor作为激励示例。

Torchdynamo 支持训练,使用 AotAutograd 来捕获反向传播:

  1. the .forward() graph and optimizer.step() is captured by torchdynamo’s python evalframe frontend

  2. for each segment of .forward() that torchdynamo captures, it uses AotAutograd to generate a backward graph segment

  3. each pair of forward, backward graph are (optionally) min-cut partitioned to save the minimal state between forward/backward

  4. the forward, backward pairs are wrapped in autograd.function modules 5. usercode calling.backward() still triggers eager’s autograd engine, which runs each ‘compiled backward’ graph as if it were one op, also running any non-compiled eager ops’ .backward() functions

您是否支持分布式代码?

DDP 已经过测试并可以正常工作,其他分布式训练库的支持正在讨论中。

使用 dynamo 进行分布式代码编写之所以具有挑战性, 是因为 AOTAutograd 会展开前向和反向传播过程, 并为后端优化提供两个图。这对于分布式代码来说是个问题, 因为我们理想情况下希望将通信操作与计算重叠进行。 Eager PyTorch 通过不同的方式(如 DDP/FSDP)实现了这一点, 使用了 autograd 钩子、模块钩子以及对模块状态的修改/变更。 在对 dynamo 的简单应用中,本应在反向传播操作期间直接执行的钩子, 可能会因为 AOTAutograd 编译函数与调度器钩子的交互方式, 被延迟到整个编译区域的反向传播操作完成之后才执行。

使用Dynamo优化DDP的基本策略概述在 distributed.py 中,主要思路是在 DDP bucket 边界 处进行图分割。

当DDP中的每个节点需要与其他节点同步其权重时,它会将其梯度和参数组织成桶,这减少了通信次数,并允许一个节点将一部分梯度广播给其他等待的节点。

分布式代码中的图中断意味着你可以期望 dynamo 及其后端优化分布式程序的计算开销,但不会优化其通信开销。如果减少的图大小剥夺了编译器的融合机会,图中断可能会影响编译速度的提升。然而,随着图规模的增加,收益会逐渐减少,因为目前大多数计算优化都是局部融合。因此,在实践中这种方法可能是足够的。

我还需要导出整个图吗?

对于绝大多数模型,你可能不需要这样做,可以使用 torch._dynamo() 优化即可,但在以下几种情况下需要完整的图,并且你可以通过简单地运行 torch.dynamo(..., nopython=True) 来确保完整的图 * 大规模训练运行,比如 $250K+,需要流水线并行和其他高级分片策略 * 推理优化器,如 TensorRTAITemplate,它们比训练优化器更激进地依赖融合操作 * 移动端训练或推理。

未来的工作将包括将通信操作追踪到图中, 协调这些操作与计算优化,并优化 通信操作。

为什么我的代码会崩溃?

如果你的代码在没有启用 dynamo 的情况下运行良好,但在启用后开始崩溃,那么最重要的第一步是确定失败发生在堆栈的哪个部分。因此,请按以下顺序运行各项操作,并且只有在上一步成功的情况下才尝试下一步。

  1. torch.compile(..., backend="eager") 仅运行 torchdynamo 正向图 捕获,然后使用 PyTorch 运行捕获的图。如果这失败了 那么就是 TorchDynamo 的问题。

  2. torch.compile(..., backend="aot_eager") 该过程会运行torchdynamo来捕获前向图,然后使用AOTAutograd 来追踪反向图,而无需任何额外的后端编译步骤。PyTorch eager将用于运行前向和反向 图。如果这失败了,则说明AOTAutograd存在问题。

  3. torch.compile(..., backend="inductor") 该步骤运行 torchdynamo 来捕获一个前向图,然后使用 AOTAutograd 和 TorchInductor 编译器来追踪反向图。如果这一步失败,则说明 TorchInductor 存在问题。

Torchdynamo 错误

如果生成的错误与"eager"后端相关,则 torchdynamo很可能是错误的来源。

要调试这些问题,我们建议将 torch._dynamo.config.verbose=True 设置为获取完整的堆栈跟踪,以便同时查看 torchdynamo 中的错误和用户代码中的错误。除了这个标志外, 你还可以通过 torch._dynamo.config.log_level 设置 torchdynamo 的 log_level。可用的日志级别如下: - logging.DEBUG: 除所有下级日志级别外,打印遇到的每条指令 - logging.INFO: 除所有下级日志级别外,打印每个被编译的函数(原始和修改后的字节码)以及捕获的图 - logging.WARNING (默认): 除所有下级日志级别外,打印图中断 - logging.ERROR: 仅打印错误

如果模型足够大,日志可能会变得难以处理。如果在模型的 Python 代码深处发生错误,仅执行发生错误的帧可能会有助于更轻松地调试。有 2 个工具可用于实现这一点:

  • env TORCHDYNAMO_DEBUG_FUNCTION=<desired_function_name> 将仅在具有该名称的函数上运行 TorchDynamo。

  • env torch._dynamo.config.replay_record_enabled = True) 在遇到错误时会转储执行记录。然后可以重放此记录,仅运行发生错误的帧。

TorchInductor 错误

使用 TorchInductor 作为选定的后端时,AOTAutograd 会用于从由 torchdynamo 捕获的前向图中生成反向图。需要注意的是,在此追踪过程中以及 TorchInductor 将前向和反向图降低为 GPU 代码或 C++ 的过程中都可能出现错误。

一个模型通常可能包含数百甚至数千个 FX 节点,因此 精确确定问题发生的具体节点可能非常困难,这也是为什么我们强烈建议你使用我们的 minifier 工具 来创建可以复现你所遇到的失败情况的小型示例。我们可以 对在 AOTAutograd 层或 Inductor 层中发生的错误进行简化,你应该按照以下顺序尝试。

  1. env TORCHDYNAMO_REPRO_AFTER="aot" python your_model.py

  2. env TORCHDYNAMO_REPRO_AFTER="dynamo" python your_model.py

缩小你的错误是最快解决问题的方法。

最小化工具实际上会在由env TORCHDYNAMO_REPRO_DIR设置的位置为您创建一个repro.py,因此请确保您对该目录具有正确的访问权限。然后您可以运行python repro.py并确认是否仍然出现相同的错误。

注意

对于其他编译器(如nvfuser),流程类似,但你会利用env TORCHDYNAMO_REPRO_AFTER="dynamo" python your_model.py

为什么编译很慢?

Dynamo 编译

TorchDynamo 有一个内置的统计功能,用于收集和显示每个编译阶段所花费的时间。这些统计信息可以通过执行 torch._dynamo 后调用 torch._dynamo.utils.compile_times() 来访问。默认情况下,这将返回一个字符串表示形式,显示每个 TorchDynamo 函数按名称划分的编译时间。

Inductor 编译

TorchInductor 内置了一个统计和跟踪功能,用于显示每个编译阶段所花费的时间、输出代码、输出图可视化和 IR 导出。env TORCH_COMPILE_DEBUG=1 python repro.py. 这是一个调试工具,旨在通过输出类似这样的内容来简化对 TorchInductor 内部的调试/理解。

调试跟踪中的每个文件都可以通过 torch._inductor.config.trace.* 启用/禁用。默认情况下,配置文件和图表都是禁用的,因为它们生成起来比较昂贵。请参阅 示例调试目录输出 以获取更多示例。

过度重新编译

当 TorchDynamo 编译一个函数(或其中一部分)时,它会做出一些关于局部变量和全局变量的假设,以允许编译器优化,并将这些假设表示为运行时检查特定值的保护机制。如果任何这些保护机制失败,Dynamo 将重新编译该函数(或其中一部分),最多可达 torch._dynamo.config.cache_size_limit 次。如果你的程序达到了缓存限制,你首先需要确定是哪个保护机制失败了,以及程序的哪一部分触发了它。

The recompilation profiler automates the process of setting TorchDynamo’s cache limit to 1 and running your program under an observation-only ‘compiler’ that records the causes of any guard failures. You should be sure to run your program for at least as long (as many iterations) as you were running when you ran into trouble, and the profiler will accumulate statistics over this duration.

from torch._dynamo.utils import CompileProfiler

prof = CompileProfiler()

def my_model():
    ...

profiler_model = torch.compile(my_model, backend=prof)
profiler_model()
print(prof.report())

许多导致图中断和过度重新编译的原因将在即将支持追踪动态张量形状、 更谨慎地选择保护措施以及更好地调整启发式算法后得到修复。

为什么要在生产环境中重新编译?

在某些情况下,你可能不希望程序运行一段时间后仍然发生意外的编译。例如,在一个对延迟要求非常严格的应用中,你正在处理生产流量。为此,TorchDynamo 提供了一种替代模式,该模式使用先前编译的图,但不会生成新的图:

frozen_toy_example = dynamo.run(toy_example)
frozen_toy_example(torch.randn(10), torch.randn(10))

你如何加速我的代码?

加速 PyTorch 代码有 3 种主要方式:

  1. 通过垂直融合实现内核融合,将连续的操作合并以避免过多的读写操作。例如,将两个连续的余弦操作融合,意味着你可以用 1 次读取和 1 次写入代替原本的 2 次读取和 2 次写入。水平融合:最简单的例子是批处理,其中单个矩阵与一批示例相乘;更一般的情况是分组 GEMM(通用矩阵乘法),其中一组矩阵乘法被一起调度。

  2. 乱序执行:一种通用的编译器优化方法。通过查看图中精确的数据依赖关系,我们可以决定执行某个节点的最佳时机,以及哪些缓冲区可以被重复使用。

  3. 自动工作分配:类似于乱序执行点, 但通过将图中的节点匹配到资源,如物理硬件或 内存,我们可以设计出合适的调度方案。

以上是加速 PyTorch 代码的一般原则,但 不同的后端会在优化内容上做出不同的权衡。例如,Inductor 首先会尽可能地进行融合操作,然后再生成 Triton 内核。它还可以

Triton 此外还提供了加速,这是由于自动内存合并、内存管理以及每个流式多处理器内的调度功能,并且其设计旨在处理分块计算。

然而,无论使用哪种后端,最好使用基准测试来查看效果,因此请尝试使用 PyTorch 性能分析器,直观检查生成的内核,并亲自观察发生了什么。

为什么我没有看到加速效果?

图中断点

你使用 dynamo 时看不到预期的加速效果,主要原因在于出现了过多的图中断。那么什么是图中断?

给定如下程序:

def some_fun(x):
    ...

torch.compile(some_fun)(x)
...

Torchdynamo 将尝试将 some_fun() 中的所有 torch/tensor 操作编译成一个 FX 图,但它可能无法将所有内容捕获到一个图中。

某些图中断的原因对 TorchDynamo 来说是无法克服的,例如调用非 torch 的 C 扩展,这些对 torchdynamo 是不可见的,可能会执行任意操作,而 TorchDynamo 无法引入必要的保护措施来确保编译后的程序可以安全地重复使用。

To maximize performance, it’s important to have as few graph breaks as possible.

识别图表中断的原因

要识别程序中的所有图中断及其相关中断原因,可以使用 torch._dynamo.explain。此工具在提供的函数上运行 TorchDynamo,并汇总遇到的图中断。以下是一个使用示例:

import torch
import torch._dynamo as dynamo
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    print("woo")
    if b.sum() < 0:
        b = b * -1
    return x * b
explanation, out_guards, graphs, ops_per_graph = dynamo.explain(toy_example, torch.randn(10), torch.randn(10))
print(explanation)
"""
Dynamo produced 3 graphs, with 2 graph break and 6 ops.
 Break reasons:
1. call_function BuiltinVariable(print) [ConstantVariable(str)] {}
   File "t2.py", line 16, in toy_example
    print("woo")

2. generic_jump
   File "t2.py", line 17, in toy_example
    if b.sum() < 0:
 """

要在此遇到第一个图中断时抛出错误,可以使用 通过使用 nopython=True 禁用 Python 回退,如果你使用过基于导出的编译器,这应该很熟悉。

def toy_example(a, b):
   ...

torch.compile(toy_example, fullgraph=True, backend=<compiler>)

当我修改代码时,为什么没有重新编译?

如果你已经通过 env TORCHDYNAMO_DYNAMIC_SHAPES=1 python model.py 启用了动态形状功能,那么在形状发生变化时,你的代码将不会重新编译。我们已添加对动态形状的支持,当形状变化小于 2 倍时,可以避免重新编译。这在像计算机视觉中图像大小不一或自然语言处理中序列长度可变等场景中特别有用。在推理场景中,通常无法提前知道批处理大小,因为你只能从不同的客户端应用程序中获取可用的数据。

一般来说,TorchDynamo 会非常努力地避免不必要的重新编译。例如,如果 torchdynamo 发现了 3 个图,而你的修改只影响了其中一个图,那么只有这个图会被重新编译。因此,另一个避免潜在缓慢编译时间的技巧是通过先编译一次模型来进行预热,之后的编译将会快得多。冷启动编译时间仍然是我们可见跟踪的一个指标。

为什么我得到的是错误的结果?

精度问题也可以通过设置环境变量 TORCHDYNAMO_REPRO_LEVEL=4来最小化,它的工作方式类似于git bisect 模型,一个完整的重现可能类似于 TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4的原因 我们需要这个是因为下游编译器会为代码生成代码,无论它是 Triton代码还是C++后端,这些下游编译器的数值计算可能会有细微的不同,但会对 你的训练稳定性产生显著影响。因此,精度调试器对我们检测代码生成中的错误或后端编译器的问题非常有用。

为什么会出现OOM错误?

Dynamo 仍是一个 alpha 产品,因此存在一些导致内存溢出(OOM)的来源。如果你遇到内存溢出问题,请按以下顺序尝试禁用以下配置,然后在 Github 上提交一个 issue,以便我们解决根本问题: 1. 如果你使用了动态形状,请尝试禁用它们,我们默认已禁用: env TORCHDYNAMO_DYNAMIC_SHAPES=0 python model.py 2. CUDA 图表与 Triton 默认在 inductor 中是启用的,但移除它们可能会缓解一些 OOM 问题: torch._inductor.config.triton.cudagraphs = False

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

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

查看教程

资源

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

查看资源