compiled mode功能

功能简介

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

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

使用方法

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

1. 用户自行准备基于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])

2. 改造模型脚本。
   1. 先处理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])

   2. 通过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)

3. 模型脚本改造后，运行并生成非封装func函数的缓存文件。
   1. 进入test.py所在目录，执行如下命令：

      cd /home/workspace
      python3 test.py

   2. 开启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”文件。

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

   新启进程，再次执行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

5. （可选）如需查看封装的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等）中。

接口说明

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)