编译时间缓存配置#
作者: Oguz Ulgen 和 Sam Larsen
简介#
PyTorch 编译器实现了多种缓存以减少编译延迟。本篇教程将演示如何配置 torch.compile 中缓存的各个部分。
先决条件#
在开始此秘籍之前，请确保您已具备以下条件
对
torch.compile的基本了解。请参阅PyTorch 2.4 或更高版本
Inductor 缓存设置#
大多数缓存都是内存中的，仅在同一进程内使用，且对用户透明。例外的是存储已编译 FX 图的缓存（FXGraphCache, AOTAutogradCache）。这些缓存允许 Inductor 在遇到具有相同张量输入形状（以及相同配置）的相同计算图时，跨进程边界避免重新编译。默认实现将编译后的产物存储在系统临时目录中。一项可选功能还支持通过将这些产物存储在 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 类似，此设置启用远程 AOTAutogradCache 功能。目前的实现使用 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#
此设置启用 TorchInductor 自动调优器（autotuner）的远程缓存。与远程 FX 图缓存类似，目前的实现使用 Redis。将其设置为 1 可启用缓存，任何其他值则禁用它。上述提到的相同主机/端口环境变量也适用于此缓存。
TORCHINDUCTOR_FORCE_DISABLE_CACHES#
将此值设置为 1 可禁用所有 Inductor 缓存。此设置对于诸如尝试冷启动编译时间或出于调试目的强制重新编译等任务非常有用。
结论#
在本教程中，我们学习了如何配置 PyTorch 编译器的缓存机制。此外，我们探索了允许用户根据其特定需求配置和优化这些缓存功能的各种设置和环境变量。