编译时缓存配置¶
作者： Oguz Ulgen 和 Sam Larsen
简介¶
PyTorch 编译器实现了多个缓存，用于减少编译延迟。此教程演示了如何配置 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，此设置启用远程 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 的自动调优器的远程缓存。与远程 FX 图缓存类似，目前的实现使用 Redis。将其设置为 1 表示启用缓存，其他值将禁用该功能。上述主机/端口环境变量同样适用于此缓存。
TORCHINDUCTOR_FORCE_DISABLE_CACHES¶
将此值设置为 1 可禁用所有 Inductor 缓存。此设置适用于诸如实验冷启动编译时间或出于调试目的强制重新编译等任务。
结论¶
在此教程中，我们学习了如何配置 PyTorch 编译器的缓存机制。此外，我们探讨了各种设置和环境变量，这些设置和变量使用户能够根据其特定需求配置和优化这些缓存功能。
