在torch.compile中的编译时缓存¶
创建于：2024年6月20日 | 最后更新：2024年12月6日 | 最后验证：2024年11月5日
作者: Oguz Ulgen 和 Sam Larsen
介绍¶
PyTorch Inductor 实现了多个缓存以减少编译延迟。
本教程展示了如何在 torch.compile 中配置缓存的各个部分。
电感器缓存设置¶
这些缓存大多数是内存中的，仅在同一进程内使用，并且对用户是透明的。一个例外是存储编译的FX图的缓存（FXGraphCache，AOTAutogradCache）。这些缓存允许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_，例如，/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类似，此设置启用了远程AOT AutogradCache功能。当前实现使用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¶
此设置启用Inductor自动调谐器的远程缓存。与远程FX图缓存一样，当前实现使用Redis。1启用缓存，任何其他值则禁用。上面列出的相同主机/端口环境变量适用于此缓存。
TORCHINDUCTOR_FORCE_DISABLE_CACHES¶
将此值设置为1以禁用所有Inductor缓存。此设置对于实验冷启动编译时间或强制重新编译以进行调试等任务非常有用。
结论¶
在本教程中，我们了解到PyTorch Inductor的缓存机制通过利用本地和远程缓存显著减少了编译延迟，这些缓存在后台无缝运行，无需用户干预。 此外，我们还探讨了各种设置和环境变量，使用户能够根据自己的特定需求配置和优化这些缓存功能。
