评价此页
★ ★ ★ ★ ★
torch.compile">
recipes/torch_compile_caching_tutorial

在 torch.compile 中使用编译时缓存#

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

作者： Oguz Ulgen

简介#

PyTorch 编译器提供了多种缓存机制，以减少编译延迟。本教程将详细解释这些机制，帮助用户选择最适合其用例的选项。

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

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

先决条件#

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

缓存机制#

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 的缓存机制通过利用本地和远程缓存，在后台显著减少编译延迟，而无需用户干预。

在此页面上
PyTorch 库