compiled mode功能

功能简介

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

TorchAir提供了一种模型缓存编译方案（cache_compile），通过将首次dynamo的编译结果落盘到磁盘的方式，加速torch.compile图模式的启动时间。

使用方法

为了端到端演示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])

改造模型脚本。

先处理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函数依赖原始模型逻辑，请用户根据实际场景需要自行定义。

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

import torch
import torch_npu
import 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所在目录，执行如下命令：
```
cd /home/workspace
python3 test.py
```

开启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

同时在当前工作目录下生成“.torchair_cache”文件夹（若无会新建），其存储了各封装func函数的缓存文件，比如“/prompt/compiled_module”和“/decode/compiled_module”文件。

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

新启进程，再次执行test.py脚本，开启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成图过程中模型脚本、模型结构、执行流程等相关信息，可用于问题定位分析。
接口调用示例如下：
```
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等）中。

接口说明

表1 cache_compile接口说明表
函数原型	def cache_compile(func, , config: Optional[CompilerConfig] = None, dynamic: bool = True, cache_dir: Optional[str] = None, global_rank: Optional[int] = None, tp_rank: Optional[int] = None, pp_rank: Optional[int] = None, *kwargs) -> Callable
函数功能	实现模型缓存编译功能。
参数（IN）	func	必选参数，缓存编译的method，只支持torch.nn.module的method，例如prompt、decode、warmup等。
	*	预留参数，可用于后续功能拓展。
	config	可选参数，图编译配置。 import torch_npu import torchair as tng config = tng.CompilerConfig() # 计算通信并行功能开关 config.experimental_config.cc_parallel_enable = True # 开启图单流执行功能 config.experimental_config.enable_single_stream = True ...... cached_model = torchair.inference.cache_compile(self._forward, config=config) 说明：本场景下，不支持同时配置dynamo export功能、使能RefData类型转换功能。
	dynamic	可选参数，是否按照输入动态trace。默认True。继承了PyTorch原有特性，具体请参考LINK。
	cache_dir	可选参数，缓存文件落盘的根目录。默认在.torchair_cache目录下。
	global_rank	可选参数，分布式训练时的rank，INT类型。取值范围为[0, world_size-1]，其中world_size是参与分布式训练的总进程数。一般情况下TorchAir会自动通过torch.distributed.get_rank()获取默认值。
	tp_rank	可选参数，指张量模型并行rank，INT类型，取值是global_rank中划分为TP域的rank id。
	pp_rank	可选参数，指流水线并行rank，INT类型，取值是global_rank中划分为PP域的rank id。
参数（OUT）	NA	-
返回值	返回一个Callable对象。
使用约束	缓存func只能被触发一次dynamo trace，也就是如果func在过程中会发生重编译，则会放弃缓存。说明：对于发生多次trace（guard失效）的函数，需要进行一次函数封装来使缓存生效，封装方式参见infer_test.py中prompt、decode封装。 func必须是method，即必须是module实例对象的方法，且该方法未被其他装饰器修饰。 func必须能成整图，即必须支持full graph。只支持推理模式，不支持带反向计算过程的func缓存。缓存文件与产生缓存时的模型脚本、输入数据、配置一一对应。若产生缓存文件后，修改了模型脚本（test.py）、配置等，缓存文件可能无法与修改后的脚本对应，需要手动删除缓存文件，重新缓存。

表2 readable_cache接口说明表
函数原型	def readable_cache(cache_bin, print_output=True, file=None)
函数功能	读取封装的func函数缓存文件compiled_module，并以可读文件（格式不限，如py、txt）格式呈现。
参数（IN）	cache_bin	必选参数，指定被封装func函数缓存文件的路径。例如/home/workspace/.torchair_cache/Model_dynamic_f2df0818d06118d4a83a6cacf8dc6d28/prompt/compiled_module。
	print_output	可选参数，是否打屏显示func函数缓存文件解析后的内容。 True：默认开启打屏显示。 False：不开启打屏显示。
参数（OUT）	file	可选参数，解析生成的可读文件路径。默认为None，即不生成可读文件。绝对路径：如/home/workspace/prompt.py 相对路径：如prompt.py，默认在调用该接口的脚本所在工作目录下。
返回值	返回文件内容，String类型。
使用约束	NA

父主题： 特性介绍