更新时间:2024-04-30 GMT+08:00
分享

训练业务代码适配昇腾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设备。

迁移操作步骤

  1. 在训练任务启动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导入

  2. “一键迁移”完成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训练入口迁移

    “一键迁移”脚本适合没有使用CUDA高阶能力的简单场景,如果涉及自定义算子、主动申请GPU显存等操作则需要额外自行适配。

  3. 手动迁移。

    手动迁移意味着需要将GPU相关的代码调用修改为NPU对应的接口,适配的内容包含但不局限于以下内容,可能会涉及到逐步的排错和适配。

    • 迁移GPU相关从cuda API调用到npu的API。

      迁移前:

      model.cuda(args.gpu)
      torch.cuda.set_device(args.gpu)
      torch.cuda.is_available()

      迁移后:

      model.npu(args.gpu)
      torch_npu.npu.set_device(args.gpu)
      torch_npu.npu.is_available()
    • 多卡多机场景下,迁移nccl相关调用到hccl。

      迁移前:

      dist.init_process_group(backend='nccl',init_method = "tcp//:127.0.0.1:**", ...... ,rank = args.rank)

      迁移后:

      dist.init_process_group(backend='hccl',init_method = "tcp//:127.0.0.1:**", ...... ,rank = args.rank)

      手动迁移可以在一键迁移有问题后,以此为线索来进行逐个适配,或者对于代码实现比较了解的前提下,主动完成代码改造。

常见问题

  1. 如何检测当前的torch_npu是否正确安装?

    可以用如下的python命令在对应的运行环境中初步校验torch_npu是否正常安装。

    python3 -c "import torch;import torch_npu;print(torch_npu.npu.is_available())"
  2. torch_npu使用报错看不懂怎么办?应该怎么求助?

    如果报错可以首先在昇腾社区的常见问题,以及Gitee的PyTorch社区中查看是否有类似的问题找到一些线索。如果还无法解决可以通过提交工单的形式从华为云ModelArts入口来进行咨询以及求助对应的专业服务。

  3. “一键迁移”似乎还是要改很多脚本才能运行起来?

    因为一键迁移其实是对于torch运行环境中常用的GPU上的接口进行和昇腾设备相应的映射,原有的训练任务代码逻辑中例如数据集导入,预训练权重的内容,以及对应的环境的超参数等内容都需要在实际的昇腾环境中进行调整。

分享:

    相关文档

    相关产品