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