Ascend-pytorch插件介绍及模型迁移
用于昇腾适配PyTorch框架，为使用PyTorch框架的开发者提供昇腾AI处理器的超强算力。
当前(2023.9.20)AscendPyTorch支持的pytorch版本如下：
| AscendPyTorch版本 | CANN版本 | 支持PyTorch版本 | 代码分支名称 |
|---|---|---|---|
| 5.0.RC2 | CANN 6.3.RC2 | 2.0.1.rc1 | v2.0.1-5.0.rc2 |
| 5.0.RC2 | CANN 6.3.RC2 | 1.8.1.post2 | v1.8.1-5.0.rc2 |
| 5.0.RC2 | CANN 6.3.RC2 | 1.11.0.post1 | v1.11.0-5.0.rc2 |
基本介绍
PyTorch Adapter实现原理
pytorch仓的结构：
- torch/ # 存放大家熟悉的python模块
- csrc # 用C++实现的pytorch前端的功能，包括python和C++的binding，autograd和JIT
- aten # "A Tensor Library"的缩写
- src/ATen # 用C++实现的tensor基本操作
- c10 # Caffe2和 A Tensor的双关语，包含pytorch的核心抽象，以及tensor和存储数据结构的实际实现
核心机制-dispatch
简单而言，就是根据API调用时，输入的数据类型来决定后端调用的API类型。比如CPU和GPU的API是不一样的，可以自动根据传入的tensor类型来自动选择API。
具体来看，对于每一个前端的算子，dispatcher会维护一个函数指针表，为每个dispatch key提供对应的视线。这个表中有针对不同后端（CPU，GPU，XLA）的dispatch条目，也有想autograd和tracing这样的高抽象层级概念的条目。dispatcher根据输入的tensor和其他东西计算出一个dispatch key，然后跳转到函数指针表所指向的函数。
所以，对于昇腾处理器而言，实现Adapter主要就是要讲用昇腾实现的算子注册到dispatcher上面，即可复用pytorch的dispatch机制完成算子分布。
示例-单算子调用流程
- 用户在前端调用算子，比如可调用nn.Module,nn.Funtional,Tensor对象上的函数；
- pybind11根据注册绑定的映射规则，调用后端C++方法；
- 后端C++接口根据输入参数来选取所需调用的算子类型（dispatch机制），比如是调用CPU实现的算子还是GPU实现的算子（注意：此处只要注册NPU实现的算子，便可调用昇腾处理器的计算能力；
- 调用相应算子，返回结果；
PyTorch Adapter的逻辑架构图
在线适配方案：模型执行，训练等主要功能流程有Pytorch框架提供，用户界面API保持不变，将Davinci设备和计算库作为扩展资源注册到PyTorch框架中。
- 优点：
- 继承PyTorch动态图特性
- 继承原生PyTorch使用方式，移植的时候，在开发方式和代码复用方便做到最小的改动；
- 继承Pytorch的原生体系结构，保留框架本身出色的特性，比如自动微分，动态分发，Debug，Profiling，Storage共享机制等；
- 扩展性：对于新增网络类型或结构，只需增加涉及的算子开发和实现。框架类算子，反向图建立和实现机制等结构可保持复用；
模型迁移
模型迁移分为以下几步：
- 准备Ascend硬件环境
- 算子满足度分析
- 模型迁移
- 自动迁移
- 手动迁移
- 模型调优
昇腾上Ascend_Pytorch安装指南
算子满足度分析
方式一：利用torch的profiler来提取模型所用到的算子，然后查看NPU对这些算子的支持情况；
方式二：直接调用CANN分析迁移工具
```bash
python ms_fmk_transplt.py -i model_file.py -o model_file_out.py
```
其中ms_fmk_transplt.py已经包含在CANN安装包中了，默认路径为：/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/ms_fmk_transplt.py。另外，该脚步还会输出迁移报告，里面有详细的不支持算子列表。
模型迁移
主要分：
- 手工迁移：用户自行修改pytorch训练脚步
- 自动迁移：
- 训练前：通过迁移脚步转换工具，自动将脚本由GPU转换为NPU版本，同时也会生成迁移报告；
- 训练时：在训练脚步中导入脚本转换库，运行训练时，自动将Pytorch训练脚本的cuda接口进行替换，操作简洁。
手工迁移
主要修改点：
- 导入NPU相关库
if torch.__version__ > 1.5:
import torch_npu
-
指定device类型为npu
device = torch.device("cuda:0") => device = torch.device("npu:0") torch.cuda.set_device(device) => torch.npu.set_device(device) tensor.to('cuda:0') => tensor.to('npu:0') -
将训练脚步中的cuda接口替换为npu接口
torch.cuda.xxx() => torch.npu.xxx() torch.cuda.set_device(0) => torch.npu.set_device(0) torch.cuda.synchronize() => torch.npu.synchronize() -
分布式代码迁移:将
nccl改为hccl
dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:29688', world_size=8, rank=rank)
# 替换为：
dist.init_process_group(backend='hccl', init_method='tcp://127.0.0.1:29688', world_size=8, rank=rank)
注：Ascend pytorch只支持DDP，不支持DP。
自动迁移
- 使用转换脚本
ms_fmk_transplt.py
python ms_fmk_transplt.py -i model_file.py -o model_file_out.py
-
导入脚步转换库(支持PT1.8以上版本，推荐这种方式)
from torch_npu.contrib import transfer_to_npu -
关于混合精度
- pytorch默认是fp32精度，而昇腾芯片的加速单元主要是fp16的算力，所以需要采用混合精度训练的方式。当前Ascend-Torch支持torch.cuda.amp和apex.amp两种方式的混合精度训练。
# 方式一：torch.cuda.amp
use_amp = True
net = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler(anabled_use_amp)
for epoch in range(epochs):
for input, target in zip(data, targets):
with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
output = net(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
opt.zero_grad()
# 方式二：apex.amp
model = torch.nn.Linear(D_in, D_out).cuda()
optimzier = torch.optim.SGD(model.parameters(), lr=1e-3)
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
...
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
...
