功能简介

使用torch.compile成图时首次编译的时间通常较长，尤其是大模型推理场景，耗时较为明显。如果遇到服务弹性扩容等业务场景，业务启动时间变长，可能会导致服务错过业务流量的高峰期。

图编译通常包括两段耗时，一段是dynamo的编译耗时，一段是GE的编译耗时。TorchAir提供了一种模型缓存编译方案（通过cache_compile接口实现），可以将首次编译结果落盘到磁盘，以加速torch.compile图模式的启动时间。

使用约束

面向推理服务场景。
该功能不支持同时配置dynamo导图功能、使能RefData类型转换功能。若开启ge_cache=True，与固定权重类输入地址功能也不能同时开启。

使用方法

为了端到端演示compiled mode功能，本章提供一个简化版的模型缓存编译使用示例，同时也展示了缓存针对特殊类型输入的处理能力（如Python Class类型、List类型等）。

准备PyTorch模型脚本。

假设在/home/workspace目录下定义了test.py模型脚本，代码示例如下：

import torch

# InputMeta为仿照VLLM(Versatile Large Language Model)框架的入参结构
@dataclasses.dataclass
class InputMeta:
    data: torch.Tensor
    is_prompt: bool

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(2, 1)
        self.linear2 = torch.nn.Linear(2, 1)
        for param in self.parameters():
            torch.nn.init.ones_(param)

    @torch.inference_mode()
    def forward(self, x: InputMeta, kv: List[torch.Tensor]):
        return self.linear2(x.data) + self.linear2(kv[0])

改造PyTorch模型脚本。

先处理forward函数。

将test.py中“forward”函数的实现提取为“_forward”函数，避免@torch.inference_mode的影响，结果如下。

@torch.inference_mode()
def forward(self, x: InputMeta, kv: List[torch.Tensor]):
    return self._forward(x, kv)
def _forward(self, x, kv):
    return self.linear2(x.data) + self.linear2(kv[0])

通过cache_compile接口实现缓存编译。

“_forward”函数是可以缓存编译的函数，但由于其会触发多次重新编译，所以要为每个场景封装一个新的func函数，然后func函数直接调用_forward函数即可。同时，在forward函数中添加调用新函数的判断逻辑。如何封装新的func函数依赖原始模型逻辑，请用户根据实际场景自行定义。

缓存func函数只能被触发一次dynamo trace，换言之如果func在过程中发生重编译，则会放弃缓存。
对于发生多次trace（guard失效）的函数，需要进行一次函数封装来使缓存生效。
func必须是method，即必须是module实例对象的方法，且该方法未被其他装饰器修饰。
func必须能成整图，即必须支持full graph。
只支持推理模式，不支持带反向计算过程的func缓存。

test.py中只展示了prompt和decode的func函数封装，具体代码示例如下：

import torch, torch_npu, torchair

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(2, 1)
        self.linear2 = torch.nn.Linear(2, 1)
        for param in self.parameters():
            torch.nn.init.ones_(param)

        # 通过torchair.inference.cache_compile实现缓存编译
        self.cached_prompt = torchair.inference.cache_compile(self.prompt)
        self.cached_decode = torchair.inference.cache_compile(self.decode)

    def forward(self, x: InputMeta, kv: List[torch.Tensor]):
        # 添加调用新函数的判断逻辑
        if x.is_prompt:
            return self.cached_prompt(x, kv)
        return self.cached_decode(x, kv)

    def _forward(self, x, kv):
        return self.linear2(x.data) + self.linear2(kv[0])
    
    # 重新封装为prompt函数
    def prompt(self, x, y):
        return self._forward(x, y)

    # 重新封装为decode函数
    def decode(self, x, y):
        return self._forward(x, y)

模型脚本改造后，运行并生成封装func函数的缓存文件。

进入test.py所在目录，执行如下命令：
1 2
cd /home/workspace python3 test.py

参考TorchAir python层日志开启INFO日志，首次执行可以看到如下屏显日志：

[INFO] TORCHAIR 2024-04-30 14:48:18 Cache ModelCacheMeta(name='CacheCompileSt.test_cache_hint.<locals>.Model.prompt(x, y)', date='2024-04-30 14:48:16.736731', version='1.0.0', fx=None) saved to /home/workspace/.torchair_cache/Model_dynamic_f2df0818d06118d4a83a6cacf8dc6d28/prompt/compiled_module
[INFO] TORCHAIR 2024-04-30 14:48:20 Cache ModelCacheMeta(name='CacheCompileSt.test_cache_hint.<locals>.Model.decode(x, y)', date='2024-04-30 14:48:19.654573', version='1.0.0', fx=None) saved to /home/workspace/.torchair_cache/Model_dynamic_f2df0818d06118d4a83a6cacf8dc6d28/decode/compiled_module

生成的各封装func函数缓存文件路径由cache_dir参数指定，缺省是当前工作路径下“.torchair_cache”文件夹（若无会新建），比如${work_dir}/.torchair_cache/${model_info}/${func}/compiled_module文件。其中${work_dir}为当前工作目录，${model_info}为模型信息，${func}为封装的func函数。

再次执行脚本，验证模型启动时间。

新启进程，再次执行test.py脚本，开启Python侧INFO日志，可以看到缓存命中的日志：

缓存文件与产生缓存时的模型脚本、输入数据、配置一一对应。若产生缓存文件后，修改了模型脚本（test.py）、配置等，缓存文件可能无法与修改后的脚本对应，需要手动删除缓存文件，重新缓存。

[INFO] TORCHAIR 2024-04-30 14:52:08 Cache ModelCacheMeta(name='CacheCompileSt.test_cache_hint.<locals>.Model.prompt(x, y)', date='2024-04-30 14:48:16.736731', version='1.0.0', fx=None) loaded from /home/workspace/.torchair_cache/Model_dynamic_f2df0818d06118d4a83a6cacf8dc6d28/prompt/compiled_module
[INFO] TORCHAIR 2024-04-30 14:52:08 Cache ModelCacheMeta(name='CacheCompileSt.test_cache_hint.<locals>.Model.decode(x, y)', date='2024-04-30 14:48:19.654573', version='1.0.0', fx=None) loaded from /home/workspace/.torchair_cache/Model_dynamic_f2df0818d06118d4a83a6cacf8dc6d28/decode/compiled_module

（可选）如需查看封装的func函数缓存文件compiled_module，通过readable_cache接口读取。

compiled_module主要存储了torch.compile成图过程中模型脚本、模型结构、执行流程等相关信息，可用于问题定位分析。
接口调用示例如下：
1 2
import torch_npu, torchair torchair.inference.readable_cache("/home/workspace/.torchair_cache/Model_dynamic_f2df0818d06118d4a83a6cacf8dc6d28/prompt/compiled_module", file="prompt.py")
compiled_module内容最终解析到可读文件prompt.py（格式不限，如py、txt等）中。

其他功能

cache_compile接口实现缓存编译时，还支持GE图编译缓存功能，通过ge_cache参数设置。

该功能通过优化GE图编译耗时，进一步加速图模式启动时间。缺省情况下，ge_cache=False（功能不开启），因受CANN包版本变更影响，一般需用户根据实际情况手动开启。

功能开启操作示例如下（以步骤2中改造脚本代码为例）：

# 开启ge_cache后的调用示例
self.cached_prompt = torchair.inference.cache_compile(self.prompt, ge_cache=True)
self.cached_decode = torchair.inference.cache_compile(self.decode, ge_cache=True)

缓存的GE编译结果文件路径与封装的func函数缓存文件路径一致，即${work_dir}/.torchair_cache/${model_info}/${func}/ge_cache_${timestamp}.om文件。其中${work_dir}为当前工作目录，${model_info}为模型信息，${func}为封装的func函数，${timestamp}为落盘的时间戳，缓存路径会自动增加ge_cache关键词。

CANN包跨版本的缓存无法保证兼容性，如果版本升级，需要清理缓存目录并重新GE编译生成缓存。
ge_cache参数不支持和固定权重类输入地址功能同时开启。
在单算子和图混跑场景下，开启该功能会额外增加一个通信域缓存，暂时不支持单算子和图同一个通信域。

模型缓存编译功能

功能简介

使用约束

使用方法

其他功能