dynamo导图功能

功能简介

离线推理场景下可通过dynamo_export接口，导出TorchAir生成的离线图（air格式）。

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

使用约束

导出时需要保证被导出部分能构成一张图。
支持单卡和多卡场景下导出图，且支持导出后带allreduce等通信类算子。
导出的air文件大小不允许超过2G（依赖的第三方库protobuf存在限制导致）。
受dynamo功能约束，不支持动态控制流if/else。
该功能暂不支持同时配置固定权重类输入地址功能。

使用方法

dynamo_export接口原型定义如下，详细的参数说明参见dynamo_export。

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

关键参数说明如下：

export_path取值：支持相对路径或绝对路径。
- 若采用相对路径：在ATC（Ascend Tensor Compiler，昇腾张量编译器）编译、执行离线模型时，需要在相对路径的父路径中执行。
- 若采用绝对路径：在ATC编译、执行离线模型时无路径限制，但是当编译好的模型拷贝至其他服务器环境时需要保证绝对路径相同，否则会找不到权重文件。

config取值：导离线图时支持config配置，具体参见表1。

表1 导图功能配置说明
支持的功能	功能说明	配置说明
auto_atc_config_generat	前端切分场景下（PyTorch模型导出时包含集合通信逻辑），是否开启自动生成ATC（Ascend Tensor Compiler，昇腾张量编译器）的json配置文件样例模板。	False（缺省值）：不开启自动生成json模板，用户自行手动配置通信域信息。 True：开启自动生成json模板。
enable_record_nn_module_stack	导出的图是否携带nn_module_stack信息，方便后端切分（PyTorch模型导出时不含集合通信逻辑，而由GE添加集合通信逻辑）运用模板。说明：前端脚本定义layer时，需要以数组的形式，即类似layer[0] = xxx，layer[1] = xxx。若不以数组形式表现变量名，相同模型结构被重复执行，从栈信息中将无法看出模型的layer结构，后端也无法切分。 record_nn_module_stack只有在model结构深度两层及以上才能获取到。	False（缺省值）：导出图不带nn_module_stack信息。 True：导出图带nn_module_stack信息。

config的配置示例如下：

        
             import torch_npu, torchair
config = torchair.CompilerConfig()
# 开启自动生成ATC的json配置文件模板
config.export.experimental.auto_atc_config_generated = True 
# 携带nn_module_stack信息
config.export.experimental.enable_record_nn_module_stack = True

导图功能支持在单卡和多卡场景下导图，且支持导出的计算图携带allreduce等通信类算子，导图结果参见产物说明，dynamo_export接口的调用示例参见使用示例。

产物说明

dynamo_export导图结果文件的缺省路径为"export_file"，用户可根据实际情况设置，产物目录如下：

└── export_file                         // 导出的文件夹，可自定义
    ├── dynamo.pbtxt                    // 导出的模型信息（可读格式）
    ├── export.air                      // 导出的模型文件（不可读），文件名缺省值为“export”
    ├── weight_xx                       // 导出的权重文件
    ├── ......                          
    ├── model_relation_config.json      // 使能auto_atc_config_generat生成的文件
    └── numa_config.json                // 使能auto_atc_config_generat生成的文件

对于导出的权重文件，其可能直接存入export.air，也可能单独存在weight_xx文件中。

当模型权重参数量不超过 (2G-200MB) 时，直接保存在export.air文件中，而权重存储路径、dtype等信息会被记录在dynamo.pbtxt的Const节点中。
当模型权重参数量超过 (2G-200MB) 时，不存入export.air文件，会在export_path路径中自动生成权重文件，（如多卡场景下dynamo export示例中p1、p2文件，文件个数取决于网络中权重的定义）。同时权重存储路径、dtype等信息会被记录在dynamo.pbtxt的FileConstant节点中。

dynamo.pbtxt ：导出的图文件，文件名固定，支持用户直接查看。该文件记录了图节点、参数数据类型/数据维度等信息。
export.air：导出的图文件，不支持用户直接查看。该文件记录了图节点、参数数据类型/数据维度等信息，某些场景下还包含模型权重信息。
weight_xx：导出的权重文件，可选产物。
model_relation_config.json ：可选产物，使能auto_atc_config_generat生成的文件，表示多个切片模型间数据关联和分布式通信组关系的配置文件。
numa_config.json：可选产物，使能auto_atc_config_generat生成的文件，用于指定目标部署环境逻辑拓扑关系的配置文件。

使用示例

单卡场景下dynamo export示例：

        
             import torch, torch_npu, torchair
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(2, 2)
        self.linear2 = torch.nn.Linear(2, 2)
        for param in self.parameters():
            torch.nn.init.ones_(param)
    def forward(self, x, y):
        return self.linear1(x) + self.linear2(y)
model = Model()
x = torch.randn(2, 2)
y = torch.randn(2, 2)
torchair.dynamo_export(x, y, model=model, dynamic=False)

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

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

多卡场景下dynamo export示例：

        
         
           
           
             import torch, os, torch_npu, 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”章节与《CANN ATC工具使用指南》中的“--cluster_config”章节。针对多卡场景，item节点被生成为node_id为0的表中，需要用户根据自己的需求手动划分至不同的node下。
mp/rank_0和mp/rank_1：生成的子目录，里面存放着每张卡的dynamo.pbtxt图信息、权重文件（若权重没有被保存在air文件中）。

父主题： 更多功能