在 torch.compile 中使用编译时缓存#
创建于:2024 年 6 月 20 日 | 最后更新:2025 年 6 月 24 日 | 最后验证:2024 年 11 月 5 日
作者: Oguz Ulgen
简介#
PyTorch 编译器提供了多种缓存机制,以减少编译延迟。本教程将详细解释这些机制,帮助用户选择最适合其用例的选项。
请参阅 编译时缓存配置,了解如何配置这些缓存。
还请查看我们的缓存基准测试:PT CacheBench 基准测试。
先决条件#
在开始此秘籍之前,请确保您已具备以下条件
基本了解
torch.compile。请参阅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 的缓存机制通过利用本地和远程缓存,在后台显著减少编译延迟,而无需用户干预。