在 torch.compile 中进行编译时缓存¶
作者：Oguz Ulgen 和 Sam Larsen
简介¶
PyTorch Inductor 实现了几种缓存来减少编译延迟。此食谱演示了如何在 torch.compile 中配置缓存的各个部分。
Inductor 缓存设置¶
大多数缓存都在内存中，仅在同一进程内使用，对用户透明。一个例外是存储已编译 FX 图的 FX 图缓存。此缓存允许 Inductor 在遇到具有相同张量输入形状（以及相同配置）的相同图时避免跨进程边界重新编译。默认实现将已编译的工件存储在系统临时目录中。一项可选功能还支持通过将这些工件存储在 Redis 数据库中来在集群内共享这些工件。
有一些与缓存相关的设置，特别是与 FX 图缓存相关的设置。这些设置可以通过以下列出的环境变量访问，或者可以在 Inductor 的配置文件中硬编码。
TORCHINDUCTOR_FX_GRAPH_CACHE¶
此设置启用本地 FX 图缓存功能，即通过将工件存储在主机的临时目录中。 1 启用，任何其他值都会禁用它。默认情况下，磁盘位置是按用户名划分的，但用户可以通过指定 TORCHINDUCTOR_CACHE_DIR（如下）来启用跨用户名的共享。
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_AUTOTUNE_REMOTE_CACHE¶
此设置为 Inductor 的自动调谐器启用远程缓存。与远程 FX 图缓存一样，当前实现使用 Redis。 1 启用缓存，任何其他值都会禁用它。上面列出的相同主机/端口环境变量适用于此缓存。
TORCHINDUCTOR_FORCE_DISABLE_CACHES¶
将此值设置为 1 以禁用所有 Inductor 缓存。此设置对于诸如尝试冷启动编译时间或强制重新编译以进行调试之类的任务很有用。
结论¶
在本教程中，我们了解到 PyTorch Inductor 的缓存机制通过利用本地和远程缓存来显著减少编译延迟，这些缓存可在后台无缝运行，无需用户干预。此外，我们还探讨了各种设置和环境变量，这些变量允许用户根据自己的特定需求来配置和优化这些缓存功能。
