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 缓存，从而实现缓存预热。

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

@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：联合图工件的缓存。

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

• AutotuningCache:

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

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

所有这些缓存工件都被写入 TORCHINDUCTOR_CACHE_DIR，默认情况下路径类似于 /tmp/torchinductor_myusername。

远程缓存

对于希望利用基于 Redis 的缓存的用户，我们还提供了远程缓存选项。请参阅 编译时缓存配置 以了解如何启用基于 Redis 的缓存。

结论

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