dynamo export功能

功能简介

TorchAir支持通过dynamo_export导出air格式的图，导出的推理模型不再依赖PyTorch框架，可直接由CANN软件栈加载执行，减少了框架调度带来的性能损耗，方便在不同的部署环境上移植。

• 导出时需要保证被导出部分能构成一张图。
• 支持单卡和多卡场景下导出图，且支持导出后带allreduce等通信类算子。
• 受dynamo功能约束，不支持动态控制流if/else。

使用方法

def dynamo_export(*args, model: torch.nn.Module, export_path: str = "export_file", export_name: str = "export", dynamic: bool = False, config=CompilerConfig(), **kwargs)

import torchair as tng
config = tng.CompilerConfig()
# 自动生成ATC的json配置文件样例
config.export.experimental.auto_atc_config_generated = True 
# 携带nn_module_stack信息
config.export.experimental.enable_record_nn_module_stack = True 

• export_path：由于导出air图时会将权重外置，权重也会被保存到该路径下，该路径同时被标记在air文件中的fileconst节点中。注意此处路径可以配置为相对路径或绝对路径，若选择相对路径，在ATC编译、执行离线模型时也需要在相对路径的父路径中执行；若选择绝对路径，在ATC编译、执行离线模型时无路径限制，但是当编译好的模型拷贝至其他服务器环境时需要保证绝对路径相同，否则会找不到权重文件。
• 对于导出的权重文件，若权重很小能被保存在export.air文件中，则不会在目录下生成额外的权重文件。若权重文件很大，导致生成的air文件过大则无法存在air文件中，权重将会转为额外的权重文件存储在导出路径中（如场景示例中的p1、p2文件）。通过查看dynamo.txt文件中是否存在FileConstant节点，确定是否生成了额外的权重文件。
• 携带nn_module_stack信息有如下约束：
  • 前端脚本定义layer时，需要以数组的形式，即类似layer[0] = xxx, layer[1] = xxx。若不以数组形式表现变量名，相同模型结构被重复执行，从栈信息中将无法看出模型的layer结构，后端也无法切分。
  • record_nn_module_stack只有在model结构深度两层及以上才能获取到。

场景示例

• 单卡场景下dynamo export示例：

  import torch
  import torch_npu
  import torchair
  
  class Model(torch.nn.Module):
  
      def __init__(self):
          super().__init__()
          self.p1 = torch.nn.Parameter(torch.randn(2, 4))
          self.p2 = torch.nn.Parameter(torch.randn(2, 4))
  
      def forward(self, x, y):
          x = x + y + self.p1 + self.p2
          return x
  
  model = Model()
  x = torch.randn(2, 4)
  y = torch.randn(2, 4)
  
  torchair.dynamo_export(x, y, model=model, export_path="./test_export_file_False", dynamic=False)

  执行如下命令查看导出结果：

  [root@localhost example_export]# tree
  .
  ├── example1.py
  └── test_export_file_False  // 指定导出的文件夹，当文件夹不存在时会自动创建
      ├── dynamo.pbtxt       // 导出可读的图信息，用于debug
      ├── export.air        // 导出的模型文件，ATC编译时的输入。其中通过fileconst节点记录了权重所在的路径与文件名
      ├── p1               // 导出的权重文件
      └── p2              // 导出的权重文件
  
  1 directory, 5 files

• 多卡场景下dynamo export示例：

  import torch, os
  import torch_npu
  import torchair
  from torchair import CompilerConfig
  
  class AllReduceSingeGroup(torch.nn.Module):
      def __init__(self):
          super().__init__()
          self.p1 = torch.nn.Parameter(torch.tensor([[1.1, 1.1], [1.1, 1.1]]))
          self.p2 = torch.nn.Parameter(torch.tensor([[2.2, 2.2], [3.3, 3.3]]))
  
      def forward(self, x, y):
          x = x + y + self.p1 + self.p2
          torch.distributed.all_reduce(x)
          return x
  
  def example(rank, world_size):
         torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
         x = torch.ones([2, 2], dtype=torch.int32)
         y = torch.ones([2, 2], dtype=torch.int32)
         mod = AllReduceSingeGroup()
         config = CompilerConfig()
         config.export.experimental.auto_atc_config_generated = True
         config.export.experimental.enable_record_nn_module_stack = True
         torchair.dynamo_export(x, y, model=mod, dynamic=True, export_path="./mp", export_name="mp_rank",
                                config=config)
  
  def mp():
      world_size = 2
      torch.multiprocessing.spawn(example, args=(world_size, ), nprocs=world_size, join=True)
  
  if __name__ == '__main__':
       os.environ["MASTER_ADDR"] = "localhost"
       os.environ["MASTER_PORT"] = "29505"
       mp()

  执行如下命令查看导出结果：

  [root@localhost example_export]# tree
  .
  ├── example1.py
  ├── example2.py
  └── mp                                       
      ├── model_relation_config.json          
      ├── mp_rank0.air                  // 第一张卡导出的模型文件
      ├── mp_rank1.air                  // 第二张卡导出的模型文件
      ├── numa_config.json          
      ├── rank_0                    // 第一张卡子目录
      │   ├── dynamo.pbtxt        // 导出可读的图信息，用于debug
      │   ├── p1                 // 导出的权重文件
      │   └── p2                 // 导出的权重文件
      └── rank_1
          ├── dynamo.pbtxt
          ├── p1
          └── p2
  
  3 directories, 12 files

  • mp：指定导出的文件夹，即export_path，当文件夹不存在时会自动创建。
  • mp_rank0/1：由指定的export_name加上rank id拼接而成。
  • model_relation_config.json、numa_config.json：前端切分场景下，自动生成ATC编译的json配置文件模板。

    json中相关字段需要用户根据自己的需求修改，字段参数含义请参考《CANN ATC工具使用指南》中“参数说明>基础功能参数>输入选项”的--model_relation_config与--cluster_config章节。针对多卡场景，item节点被生成为node_id为0的表中，需要用户根据自己的需求手动划分至不同的node下。

  • mp/rank_0和mp/rank_1：生成的子目录，里面存放着每张卡的dynamo.pbtxt图信息、权重文件（若权重没有被保存在air文件中）。