torch.compile 中的编译时缓存¶
创建于：2024 年 6 月 20 日 | 最后更新：2024 年 12 月 06 日 | 最后验证：2024 年 11 月 05 日
作者： Oguz Ulgen 和 Sam Larsen
简介¶
PyTorch Inductor 实现了多个缓存以减少编译延迟。本食谱演示了如何在 torch.compile 中配置缓存的各个部分。
Inductor 缓存设置¶
大多数这些缓存都是内存中的，仅在同一进程内使用，并且对用户是透明的。例外情况是存储编译后的 FX 图的缓存（FXGraphCache、AOTAutogradCache）。这些缓存允许 Inductor 在遇到具有相同 Tensor 输入形状（和相同配置）的相同图时，避免跨进程边界重新编译。默认实现将编译后的工件存储在系统临时目录中。一个可选功能还支持通过将这些工件存储在 Redis 数据库中，在集群内共享它们。
有一些设置与缓存，特别是 FX 图缓存相关。这些设置可以通过下面列出的环境变量访问，也可以在 Inductor 的配置文件中硬编码。
TORCHINDUCTOR_FX_GRAPH_CACHE¶
此设置启用本地 FX 图缓存功能，即通过将工件存储在主机的临时目录中。1 启用，任何其他值禁用它。默认情况下，磁盘位置是按用户名划分的，但用户可以通过指定 TORCHINDUCTOR_CACHE_DIR（如下所示）来启用跨用户名共享。
TORCHINDUCTOR_AUTOGRAD_CACHE¶
此设置扩展了 FXGraphCache，以在 AOTAutograd 级别而不是 Inductor 级别存储缓存结果。1 启用，任何其他值禁用它。默认情况下，磁盘位置是按用户名划分的，但用户可以通过指定 TORCHINDUCTOR_CACHE_DIR（如下所示）来启用跨用户名共享。TORCHINDUCTOR_AUTOGRAD_CACHE 需要 TORCHINDUCTOR_FX_GRAPH_CACHE 才能工作。相同的缓存目录存储 AOTAutogradCache（在 {TORCHINDUCTOR_CACHE_DIR}/aotautograd 下）和 FXGraphCache（在 {TORCHINDUCTOR_CACHE_DIR}/fxgraph 下）的缓存条目。
TORCHINDUCTOR_CACHE_DIR¶
此设置指定所有磁盘缓存的位置。默认情况下，该位置位于系统临时目录下的 torchinductor_<username> 中，例如 /tmp/torchinductor_myusername。
请注意，如果未在环境中设置 TRITON_CACHE_DIR，则 Inductor 会将 Triton 缓存目录设置为相同的临时位置，位于 Triton 子目录下。
TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE¶
此设置启用远程 FX 图缓存功能。当前实现使用 Redis。1 启用缓存，任何其他值禁用它。以下环境变量配置 Redis 服务器的主机和端口
TORCHINDUCTOR_REDIS_HOST（默认为 localhost） TORCHINDUCTOR_REDIS_PORT（默认为 6379）
请注意，如果 Inductor 找到远程缓存条目，它会将编译后的工件存储在本地磁盘缓存中；该本地工件将在同一台机器上的后续运行中提供。
TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE¶
与 TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE 类似，此设置启用远程 AOT AutogradCache 功能。当前实现使用 Redis。1 启用缓存，任何其他值禁用它。以下环境变量配置 Redis 服务器的主机和端口：TORCHINDUCTOR_REDIS_HOST（默认为 localhost） TORCHINDUCTOR_REDIS_PORT（默认为 6379）
TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE` 依赖于 TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE 才能工作。相同的 Redis 服务器可以存储 AOTAutograd 和 FXGraph 缓存结果。
TORCHINDUCTOR_AUTOTUNE_REMOTE_CACHE¶
此设置启用 Inductor 自动调谐器的远程缓存。与远程 FX 图缓存一样，当前实现使用 Redis。1 启用缓存，任何其他值禁用它。上面列出的相同主机/端口环境变量适用于此缓存。
TORCHINDUCTOR_FORCE_DISABLE_CACHES¶
将此值设置为 1 以禁用所有 Inductor 缓存。此设置对于诸如实验冷启动编译时间或强制重新编译以进行调试等任务非常有用。
结论¶
在本食谱中，我们了解到 PyTorch Inductor 的缓存机制通过利用本地和远程缓存显着减少了编译延迟，这些缓存机制在后台无缝运行，无需用户干预。此外，我们还探讨了各种设置和环境变量，用户可以使用这些设置和环境变量根据其特定需求配置和优化这些缓存功能。
