评价此页
torch.compile">

torch.compile 中的编译时缓存#

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

作者: Oguz Ulgen

简介#

PyTorch 编译器提供多种缓存方案以减少编译延迟。本篇教程将详细解释这些方案,帮助用户为自己的用例选择最佳选项。

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

另请参阅我们的缓存基准测试:PT CacheBench 基准测试

先决条件#

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

缓存方案#

torch.compile 提供以下缓存方案

  • 端到端缓存(也称为 Mega-Cache

  • TorchDynamoTorchInductorTriton 的模块化缓存

请注意,缓存机制会验证缓存工件是否与当前的 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)

此操作将填充下一节中讨论的所有模块化缓存,包括 PGOAOTAutogradInductorTritonAutotuning

TorchDynamoTorchInductorTriton 的模块化缓存#

上述 Mega-Cache 由多个组件构成,这些组件无需用户干预即可使用。默认情况下,PyTorch 编译器为 TorchDynamoTorchInductorTriton 提供本地磁盘缓存。这些缓存包括

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

  • TritonCache:Triton 编译结果的缓存,包括由 Triton 生成的 cubin 文件和其他缓存工件。

  • InductorCacheFXGraphCacheTriton 缓存的集合。

  • AOTAutogradCache:联合图工件的缓存。

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

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

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

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

远程缓存#

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

结论#

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