深度学习训练业务代码适配昇腾PyTorch代码适配
训练业务代码适配昇腾PyTorch代码适配
前提条件
- 要迁移的训练任务代码在GPU上多次训练稳定可收敛。训练业务代码和数据，应该确保在GPU环境中能够运行，并且训练任务有稳定的收敛效果。
- 本文只针对基于PyTorch的训练脚本迁移。这里假设用户使用的是基于PyTorch的训练代码进行迁移。其他的AI引擎如TensorFlow、Caffe等不在本指导的讨论范围中。
- 已经完成环境准备（参考环境准备），并且代码、预训练模型、数据等训练必需内容已经上传到环境中。
约束和限制
- 安装插件后，大部分能力能够对标在GPU上的使用，但是不是100%的行为和GPU上是一一对应的，比如在torch_npu下，一个进程只能操作一张昇腾卡，而没有一个进程可以操作多卡的能力等。
- 基于PyTorch上的第三方开发库非常多，例如transformers、accelerate、deepspeed以及Megatron等，这些三方库昇腾也做了类似PyTorch Adapter的适配插件库，可以在Gitee的昇腾官方仓库中找到，请按需进行使用。
代码迁移基础知识
- PyTorch官方并不直接支持昇腾的后端，所以官方的版本无法直接利用昇腾设备完成训练加速。当前PyTorch直接支持的后端包括CUDA和AMD ROCm。
- PyTorch Adapter作为一个PyTorch“插件”，在已安装PyTorch的基础上安装后，支持在不改变PyTorch表达层的基础上，动态添加昇腾后端适配，包含增加了NPU设备、hccl等一系列能力的支持。安装后可以直接使用PyTorch的表达层来运行在NPU设备上。
- 当前提供了“一键迁移”脚本进行GPU到昇腾适配，原理是通过monkey-patch的方式将torch下的CUDA、nccl等操作映射为NPU和hccl对应的操作。如果没有用到GPU的高阶能力，例如自定义算子、直接操作GPU显存等操作，简单场景下可以直接使用“一键迁移”。
图1 torch_npu工作原理示意图
- NPU（Neural Network Processing Unit）和GPU在构造结构上存在差异，因此迁移过程并不是完全平替的关系。昇腾训练芯片属于NPU的范畴，虽然在表达层可以通过torch.cuda和torch.npu的形式来替代，但是真实的算子下发、显存管理、集合通信等，在进阶问题分析时需要了解NPU的运行机制，能够更好的使用NPU设备。
迁移操作步骤
- 在训练任务启动Python脚本入口，初始化PyTorch Adapter（torch_npu）。
在torch_npu安装后，该部分并没有直接植入到PyTorch中生效，其实是在显式的调用后，前端通过monkey-patch的方式注入到torch对象中，后端的注册了NPU设备，以及HCCL的参数面通信能力可以直接使用。
#torch npu初始化 import torch_npu
在执行对应的导入torch_npu代码后，可以直接使用torch.npu相关接口和能力。
图2 torch_npu导入
- “一键迁移”完成GPU代码到昇腾的快速适配。
torch_npu初始化后，原则上需要用户将原来代码中CUDA相关的内容适配到NPU相关的接口上，包含算子API、显存操作、数据集操作，以及分布式训练的参数面通信nccl适配到hccl上，手动操作修改点相对较为分散，可以通过transfer_npu的模式来进行快速适配。
一键迁移的原理是，通过注入的方式将当前Python运行环境中，运行时的torch.cuda等需要适配的接口和操作都映射成为torch.npu对应的接口。所以理论上常见的场景下当前的代码不需要额外的适配就可以运行到昇腾设备上了。
#自动映射cuda API到npu的代码 from torch_npu.contrib import transfer_to_npu
图3 一键迁移后cuda映射为npu相关的API
以chatGLM-6b为示例，在使用一键迁移时，在开发环境中克隆对应的代码，假设数据和预训练权重已经配置好，可以直接在ptuning目录下，训练入口代码main.py中添加两行代码来完成昇腾运行适配，注意添加位置为导入torch之后。启动训练脚本可以观察运行效果。
图4 chatGLM-6b pTuning训练入口迁移
对于我自己代码来说就加了几行代码而已：如下try except部分
import os
import random
import torch
import time
import sys
import subprocess
try:
import torch_npu
print("torch npu available:", torch_npu.npu.is_available())
from torch_npu.contrib import transfer_to_npu
print("torch cuda available:", torch.cuda.is_available())
except Exception as e:
print("import torch_npu not available:", e)
