编译时缓存配置 ¶
作者:Oğuz Ulgen 和 Sam Larsen
简介
PyTorch 编译器实现了多个缓存以减少编译延迟。本菜谱演示了您如何配置缓存的不同部分 torch.compile
。
前提条件 _
在开始此菜谱之前,请确保您有以下条件:
了解
torch.compile
的基本概念。请参阅:PyTorch 2.4 或更高版本
电感器缓存设置 ¶
这些缓存大多在内存中,仅用于同一进程,对用户是透明的。一个例外是存储编译后的 FX 图的缓存( FXGraphCache
, AOTAutogradCache
)。这些缓存允许 Inductor 在遇到具有相同 Tensor 输入形状(以及相同配置)的相同图时,避免跨进程边界重新编译。默认实现将编译后的工件存储在系统临时目录中。可选功能还支持通过将它们存储在 Redis 数据库中来在集群内共享这些工件。
有一些与缓存相关的设置,特别是与 FX 图缓存相关。这些设置可以通过以下环境变量访问,也可以在 Inductor 的配置文件中硬编码。
TORCHINDUCTOR_FX_GRAPH_CACHE
此设置启用本地 FX 图缓存功能,该功能将工件存储在宿主机的临时目录中。将设置设置为 1
启用此功能,而任何其他值都将禁用它。默认情况下,磁盘位置按用户名划分,但用户可以通过指定 TORCHINDUCTOR_CACHE_DIR
(如下所示)来启用跨用户名的共享。@4#需要@5#才能工作。相同的缓存目录存储@6#(在@7#下)和@8#(在@9#下)的缓存条目。
火炬诱导器自动微分缓存
此设置将 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_<username>
,例如, /tmp/torchinductor_myusername
。
注意,如果未在环境中设置 TRITON_CACHE_DIR
,Inductor 会将 Triton
缓存目录设置为相同的临时位置,在 Triton 子目录下。
火炬诱导器 FX 图远程缓存
此设置启用了远程 FX 图形缓存功能。当前实现使用 Redis
。 1
启用缓存,任何其他值将禁用它。以下环境变量配置 Redis 服务器的主机和端口:
TORCHINDUCTOR_REDIS_HOST
(默认为 localhost
) TORCHINDUCTOR_REDIS_PORT
(默认为 6379
)
备注
注意,如果 Inductor 找到远程缓存条目,它将编译的工件存储在本地磁盘缓存中;该本地工件将在同一台机器上的后续运行中提供。
火炬诱导器自动微分远程缓存
与 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
自动调优的远程缓存。类似于远程 FX 图缓存,当前实现使用 Redis。将其设置为 1
启用缓存,而任何其他值将禁用它。上述提到的相同的主机/端口环境变量适用于此缓存。
火炬诱导器_强制禁用缓存
将此值设置为 1
以禁用所有电感器缓存。此设置对于实验冷启动编译时间或强制重新编译以进行调试等任务很有用。
结论 ¶
在本教程中,我们学习了如何配置 PyTorch 编译器的缓存机制。此外,我们还探讨了各种设置和环境变量,这些设置和环境变量允许用户根据其特定需求配置和优化这些缓存功能。