编译时缓存配置¶
作者： 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 编译器的缓存机制。此外，我们还探索了各种设置和环境变量，这些设置和环境变量允许用户根据其特定需求配置和优化这些缓存功能。
