torch.compile">

recipes/torch_compile_caching_tutorial

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

torch.compile 中的编译时缓存

创建日期：2024 年 6 月 20 日 | 最后更新：2025 年 6 月 24 日 | 最后验证：2024 年 11 月 5 日

作者： Oguz Ulgen

简介

PyTorch 编译器提供了多种缓存方案以减少编译延迟。本篇指南将详细介绍这些方案，帮助用户为自己的应用场景选择最佳选项。

请查阅 编译时缓存配置 以了解如何配置这些缓存。

另请查看我们的缓存基准测试：PT CacheBench 基准测试。

先决条件

在开始此秘籍之前，请确保您已具备以下条件

• 需要对 torch.compile 有基本的了解。请参阅

  • torch.compiler API 文档

  • torch.compile 简介

  • Triton 语言文档

• PyTorch 2.4 或更高版本

缓存方案

torch.compile 提供以下缓存方案

• 端到端缓存（也称为 Mega-Cache）

• TorchDynamo、TorchInductor 和 Triton 的模块化缓存

请注意，缓存机制会验证缓存构件是否在相同的 PyTorch 和 Triton 版本下使用，以及当设备设置为 cuda 时是否使用相同的 GPU。

torch.compile 端到端缓存（Mega-Cache）

端到端缓存（以下称为 Mega-Cache）是寻求可移植缓存解决方案的用户的理想选择，它可以存储在数据库中，并可在稍后（可能是在另一台机器上）进行提取。

Mega-Cache 提供两个编译器 API

• torch.compiler.save_cache_artifacts()

• torch.compiler.load_cache_artifacts()

预期的工作流程是：在编译并执行模型后，用户调用 torch.compiler.save_cache_artifacts()，该函数将以可移植的形式返回编译器构件。随后，用户可能在另一台机器上，使用这些构件调用 torch.compiler.load_cache_artifacts()，从而预先填充 torch.compile 缓存，以实现缓存加速（jump-start）。

参考以下示例。首先，编译并保存缓存构件。

@torch.compile
def fn(x, y):
    return x.sin() @ y

a = torch.rand(100, 100, dtype=dtype, device=device)
b = torch.rand(100, 100, dtype=dtype, device=device)

result = fn(a, b)

artifacts = torch.compiler.save_cache_artifacts()

assert artifacts is not None
artifact_bytes, cache_info = artifacts

# Now, potentially store artifact_bytes in a database
# You can use cache_info for logging

之后，可以通过以下方式启动缓存：

# Potentially download/fetch the artifacts from the database
torch.compiler.load_cache_artifacts(artifact_bytes)

此操作将填充下一节中讨论的所有模块化缓存，包括 PGO、AOTAutograd、Inductor、Triton 和 Autotuning。

TorchDynamo、TorchInductor 和 Triton 的模块化缓存

上述 Mega-Cache 由无需用户干预即可使用的独立组件组成。默认情况下，PyTorch 编译器带有用于 TorchDynamo、TorchInductor 和 Triton 的本地磁盘缓存。这些缓存包括：

• FXGraphCache：编译中使用的基于图的 IR 组件缓存。

• TritonCache：Triton 编译结果缓存，包括由 Triton 生成的 cubin 文件及其他缓存构件。

• InductorCache：FXGraphCache 和 Triton 缓存的集合。

• AOTAutogradCache：联合图（joint graph）构件缓存。

• PGO-cache：动态形状决策的缓存，用于减少重新编译次数。

• AutotuningCache:

  • Inductor 生成 Triton 内核并对其进行基准测试，以选择最快的内核。

  • torch.compile 的内置 AutotuningCache 会缓存这些结果。

所有这些缓存构件都会写入 TORCHINDUCTOR_CACHE_DIR，默认路径类似于 /tmp/torchinductor_myusername。

远程缓存

对于希望利用基于 Redis 缓存的用户，我们也提供远程缓存选项。查看 编译时缓存配置 以了解更多关于如何启用基于 Redis 的缓存的信息。

结论

在本指南中，我们了解到 PyTorch Inductor 的缓存机制通过利用本地和远程缓存显著减少了编译延迟，这些机制在后台无缝运行，无需用户干预。