Zhao Dongyu's Blog

A life which is unexamined is not worth living.

0%

torch.compile

Posted on Edited on In
Symbols count in article: 7.5k Reading time ≈ 7 mins. 
“`torch.compile“` speeds the flame,
Trade-offs linger, but worth the game.
Train or infer, it cuts the line,
With care and craft, its power’s thine.

开始学习 torch.compile ~

torch.compiler API documentation

  1. torch.compiler 概述

torch.compiler 是 PyTorch 2.x 中引入的一个命名空间，它提供了一些内部编译方法供用户使用。是一个旨在解决 PyTorch 中准确捕获计算图问题(accurate graph capturing)的函数，目标是帮助用户更快地运行 PyTorch 程序。

从 C++ 到 Python 的转变：torch.compile 是用 Python 编写的，标志着 PyTorch 从 C++ 向 Python 的转变。

  1. torch.compile 的底层技术 torch.compile 依赖以下关键技术：
  • TorchDynamo (torch._dynamo) : capture PyTorch graphs
    • 这是一个内部 API，利用 Python 的 Frame Evaluation API 安全地捕获 PyTorch 计算图。TorchDynamo 的一些外部方法通过 torch.compiler 命名空间暴露给用户。
  • TorchInductor : deep learning compiler
    • 这是 torch.compile 的默认深度学习编译器，能够为多种加速器和后端生成快速代码。对于 NVIDIA、AMD 和 Intel GPU，它使用 OpenAI Triton 作为关键组件。
  • AOT Autograd （Ahead-of-Time Autograd）
    • 它不仅捕获用户代码，还捕获反向传播，从而实现对前向和反向传播的同时加速。

As mentioned above, to run your workflows faster, torch.compile through TorchDynamo requires a backend that converts the captured graphs into a fast machine code. Different backends can result in various optimization gains. The default backend is called TorchInductor, also known as inductor, TorchDynamo has a list of supported backends developed by our partners, which can be see by running torch.compiler.list_backends() each of which with its optional dependencies.

Introduction to torch.compile

torch.compile makes PyTorch code run faster by JIT-compiling PyTorch code into optimized kernels, all while requiring minimal code changes.

通过将 PyTorch 代码即时编译（JIT）为优化后的内核，从而实现加速，同时只需要极少的代码更改。

简单使用 @torch.compiler 进行装饰即可

当你对目标函数或模块使用 torch.compile 时，编译器会尝试递归地编译目标函数或模块内部的每一个函数调用（除非这些函数在跳过列表中，例如内置函数或 torch.* 命名空间中的某些函数）。

可以通过使用 torch.compiler.disable 来禁用某些函数的编译。

@torch.compiler.disable(recursive=False)

Best Practices:

    1. Top-Level Compilation 顶层编译
    • 一种方法是尽可能在最高级别进行编译（例如，在顶层模块初始化/调用时），并在遇到过多的图中断或错误时选择性地禁用编译。如果仍然存在许多编译问题，可以尝试分别编译各个子组件。
    1. Modular Testing 模块化测试
    • 在将函数和模块集成到更大的模型之前，使用 torch.compile 对它们进行单独测试，以隔离潜在问题。
    1. Disable Compilation Selectively 选择性禁用编译
    • 如果某些函数或子模块无法被 torch.compile 处理，可以使用 torch.compiler.disable 上下文管理器递归地将它们从编译中排除。
    1. Compile Leaf Functions First 先编译叶子函数(Leaf Functions, 内部不调用任何其他函数的函数)

torch.compile 的首次运行时间比 eager mode 长得多。这是因为 torch.compile 在执行时会将模型编译为optimized kernels。模型结构不改变就不需要重新编译。如果多次运行优化后的模型，与 eager mode 相比，会看到显著的改进。

加速主要来自于减少 Python overhead 和 GPU read/writes

第二次运行经过 torch.compile 优化的模型时，时间明显比其他运行时间长，尽管它比第一次运行快得多。这是因为 reduce-overhead 模式会运行几次 CUDA graphs 的预热迭代 (a few warm-up iterations)。

Primarily, the advantage of torch.compile lies in its ability to handle arbitrary Python code with minimal changes to existing code.

相比于现有的 PyTorch 编译解决方案（TorchScript 或 FX Tracing），主要优势在于， torch.compile 能够以最小的代码更改量处理任意 Python 代码。

TorchDynamo 负责将任意 Python 代码即时编译（JIT）为 FX 图，这些图随后可以被进一步优化。TorchDynamo 通过在运行时分析 Python 字节码并检测对 PyTorch 操作的调用，来提取 FX 图。通常情况下，torch.compile 的另一个组件 TorchInductor 会进一步将 FX 图编译为优化后的内核。但 TorchDynamo 允许使用不同的后端。为了查看 TorchDynamo 输出的 FX 图，我们可以创建一个自定义后端，该后端输出 FX 图并简单地返回图的未优化前向方法。

TorchInductor: a PyTorch-native Compiler with Define-by-Run IR and Symbolic Shapes

TorchDynamo 是通过动态 Python 字节码转换来解决 the graph capture problem 的，

TorchInductor 是 PyTorch 的一个新编译器，它能够表示 PyTorch 的全部功能，并且以一种通用的方式构建，以便能够支持训练和多个后端目标。

其设计理念是通过一种轻薄且易于修改的方式，将 PyTorch 符号化地映射到低级后端，支持在不同后端之间快速实验和自动调优，以及更高层次的优化，例如内存规划(memory planning)。

为了迫使 TorchInductor 的设计具有通用性，从两个低级执行目标开始，它们代表了设计空间中的不同点：

  • Triton 是一种新的编程语言，其生产力远高于 CUDA，但能够以简洁易懂的代码超越高度优化的库（如 cuDNN）的性能。它由 OpenAI 的 Philippe Tillet 开发，并且在行业中获得了巨大的采用和牵引力。Triton 支持 NVIDIA GPU，并且作为一种替代手写 CUDA 内核的方案，其受欢迎程度正在迅速增长。

  • C++/OpenMP 是一种广泛采用的规范，用于编写并行内核。OpenMP 提供了一种工作共享的并行执行模型，并支持 CPU。C++ 作为一种高度可移植的语言，也很有趣，因为它可以支持导出到更奇特的边缘设备和硬件架构。

vLLM’s torch.compile integration

在vLLM V1架构中，torch.compile默认启用，是框架的关键部分。

编译缓存（Compilation Cache）

vLLM 会根据多种因素（包括模型配置、PyTorch 配置、模型的前向传播函数及其调用的相关函数）决定一个目录来存储编译产物。这些缓存可以被直接复制到部署环境中，从而节省大量的编译时间，加速 vLLM 实例的启动时间。vLLM 保证所有编译工作在服务请求之前完成。

Python 代码编译（Python Code Compilation）

通过 PyTorch 的 Dynamo 工具来捕获模型的前向传播函数及其调用的其他函数。这些函数会被编译成一个新的函数，并存储在缓存目录中。

  • 编译后的代码会生成一个计算图（ computation_graph.py ）和一个转换后的代码文件（ transformed_code.py ）。
  • 任何在这些文件中的代码更改都会触发缓存失效，从而重新编译。

计算图处理（Computation Graph Processing）

计算图中每个张量都有形状注释，输入包括输入 ID、位置 ID、模型权重和缓冲区，输出是最终的隐藏状态。

计算图的大部分输入具有静态形状（因为它们是模型权重和缓冲区），只有输入 ID 和位置 ID 具有符号形状（即形状可以在批次之间变化）。

  • attention operation 比较复杂，需要与键值缓存（kv caches）交互，但其输出与输入查询的形状相同。因此，vLLM 将整个 attention operation 封装为一个 PyTorch 自定义操作 torch.ops.vllm.unified_attention_with_output ，以便 Dynamo 不会尝试检查其内部操作。
  • 计算图被进一步拆分为多个子模块，每个子模块是一个计算图片段，可以独立处理。

计算图编译（Computation Graph Compilation）

使用 PyTorch 的 Inductor 编译器来编译计算图的各个片段。每个片段的编译结果会被存储在缓存目录中。如果缓存目录已经存在（例如第二次运行相同的代码），Inductor 编译将被完全跳过，直接从磁盘加载上次的编译产物。

为特定的形状（例如特定的 batch size ）编译内核。这种情况下，计算图中的所有形状都是静态且已知的，Inductor 会启用 自动调优（auto-tuning） 来优化性能。

Cudagraph 捕获（Cudagraph Capture）

vLLM 的 V1 架构使用分段(piecewise)的 Cudagraph 。计算图被拆分为多个片段，每个片段在 attention operations 之间（including the first graph before any attention operation, and the last graph after all the attention operation）捕获 Cudagraph。

  • 这种设计基于一个常见的观察结果：注意力操作之间的计算通常是 token-wise 的，适合 Cudagraph；而 the attention operation 本身比较复杂，不适合 Cudagraph。因此，通过以 Eager 模式 running the attention operation，而以 Cudagraph 模式运行其余操作，保持了 attention operation 的灵活性。

  • 分段的 Cudagraph 还具有细粒度 (fine-grained) 的内存管理，目的是将 attention kernel 排除在 Cudagraph 之外，同时将所有其他模块和内存分配操作保留在 Cudagraph 中。 --> 这就是为什么V1中的注意力操作将输出张量作为注意力的输入。

在 vLLM 中，Cudagraph 的捕获和管理由编译器后端负责，执行时会根据 batch size 自动选择合适的 Cudagraph。模型调用者只需要正确管理输入缓冲区，而中间缓冲区的管理由编译器后端自动完成。 这种设计允许模型在不同的batch size下灵活运行，同时保持高性能。编译器后端可以根据不同的 batch size 捕获和管理多个 Cudagraph，确保在不同场景下都能高效运行。

  • 1.Cudagraph 的捕获和管理
    • 捕获（Capture）：Cudagraph 的捕获是指将一系列操作（如模型的前向传播过程中的某些片段）记录下来，形成一个可重复执行的图。在 vLLM 中，这些操作主要是模型计算图中的一部分，特别是那些在注意力操作之间的部分。
    • 管理（Manage）：捕获的 Cudagraph 由编译器后端（compiler backend）管理。这意味着编译器后端负责存储这些图，并在需要时加载它们。编译器后端还负责确保这些图在正确的条件下被正确地执行。
    1. Cudagraph 的执行（Replay）
    • 执行条件：Cudagraph 的执行取决于当前的批量大小（batch size）。如果当前的批量大小与之前捕获的某个 Cudagraph 的批量大小匹配，那么这个 Cudagraph 就会被执行（replayed）。
    • 执行过程：当匹配的 Cudagraph 被执行时，编译器后端会自动管理所有中间缓冲区（intermediate buffers）。这意味着模型的调用者（model runner）不需要关心这些中间缓冲区的分配和释放，编译器后端会自动处理。

Ways to use torch.compile

主要探讨了 PyTorch 的 torch.compile 功能在不同场景下的使用方法、优缺点以及一些实际应用案例。

Some tips:

优点

  • 潜在性能提升：虽然不能简单地“一键加速”，但通过一些工作，几乎总能获得性能提升。

缺点 Downsides:

  • The compiler is complicated. 编译器复杂性：编译器的复杂性意味着用户需要投入一定的时间和精力来理解和使用 torch.compile 。

  • Compile time can be long. 编译时间可能较长：编译模型需要时间，这可能会影响小型实验或频繁崩溃的任务的效率。此外，如果任务在支持抢占的系统上运行，可能会导致重复编译。

  • Numerics divergence from eager. 数值结果可能与 eager mode 不同：编译器可能选择不同的矩阵乘法算法或优化半精度计算，这可能导致结果与 eager mode 不完全一致，甚至可能影响模型的收敛。

  • Autotuning makes the most sense for inference.

  • Warmup inference processes before serving traffic to them.

  • Try skip_guard_eval_unsafe to reduce guard overhead.

探讨 torch.export 的使用方式，这可能是解决一些 torch.compile 问题的另一种途径。

missing manual

Do an ablation

If PyTorch crashes, it’s usually pretty obvious what caused the crash. But if your outputs are garbage, this is kind of the nightmare scenario where there’s no where to start looking for the problem.

有两种主要的消融方法：

  • 可以禁用编译器堆栈的层级（例如，禁用 Inductor），并尝试定位问题所在。为了以这种方式进行消融，修改 torch.compile 调用中的 backend 参数。我们推荐测试以下三种设置：
    • backend="eager" ：如果这失败了，这表明是 Dynamo 的问题。
    • backend="aot_eager" ：如果这失败了，但 eager 没有失败，这表明是 AOTAutograd 的问题。
    • backend="aot_eager_decomp_partition" ：如果这失败了，但 aot_eager 没有失败，这表明问题与我们的分解/分区器有关。
    • 此外，如果使用的是 mode="reduce-overhead" ，应该尝试不使用它（如果是 CUDA 图问题）。如果使用的是 dynamic=True ，应该尝试不使用它（如果是动态形状问题）。如果怀疑问题与某个特定的 FX 传递有关，也可以手动禁用该传递，通过注释掉或禁用相关的配置，例如 torch/_inductor/fx_passes/joint_graph.py 或 torch/_inductor/fx_passes/post_grad.py （TODO：优化燃料，这样就不需要了解内部机制来做到这一点）。
  • 可以禁用模型层级的编译器。可以通过将 torch.compile 调用推入模型内部来手动完成，或者使用类似 https://gist.github.com/ezyang/4a5138b11327335e618dd37ad2fd0a4e 的补丁程序来程序化地完成（TODO：正式落地）。

还有其他方法可以诊断准确性问题（例如，逐层比较输出），但消融实验很容易进行，且不费力，因此应该先尝试它们。

使用缓存加速编译 Speeding up compilation with caching

默认情况下，我们也有一个文件系统缓存，保存在 /tmp 中。如果系统中 /tmp 的保留时间不够长，可以使用 TORCHINDUCTOR_CACHE_DIR 更改缓存目录。

使用分层编译加速编译 Speeding up compilation with hierarchical compilation

默认情况下，PT2 将您的所有模型代码内联到一个函数中，然后对其进行编译。在某些架构中，这意味着重复使用的块（例如 Transformer 块）将被内联并反复编译。如果您没有从跨块边界的融合中受益，您可以通过仅对块本身使用 torch.compile 来显著减少编译时间，这样它只会被编译一次，并在每个实例中重复使用。您可能需要设置 torch._dynamo.config.inline_inbuilt_nn_modules = True 以确保在 self 实例发生变化时不会重新编译。

Thanks for your support.