Pytorch的一小步，昇腾芯片的一大步

Pytorch的一小步，昇腾芯片的一大步

相信在AI圈的人多多少少都看到了最近的信息：PyTorch最新2.1版本宣布支持华为昇腾芯片！

1、 发生了什么事儿？

在2023年10月4日PyTorch 2.1版本的发布博客上，PyTorch介绍的beta版本新特性上有一个PRIVATEUSE1特性是提高了第三方设备的支持，并说明了华为Ascend NPU（昇腾NPU芯片）的OSS小组已经成功将torch_npu整合进入PyTorch;

人话就是： pytorch原生支持的是CPU和CUDA，现在昇腾NPU也可以支持pytorch了；

• **PyTorch2.1 发布了PRIVATEUSE1新特性，PyTorch介绍的beta版本PRIVATEUSE1新特性，**主要目标是让PyTorch可以使用更多的硬件；

• PyTorch 2.1对华为昇腾芯片NPU的支持：简单理解就是华为的Ascend NPU（昇腾NPU芯片）的OSS小组通过PyTorch的PRIVATEUSE1特性让PyTorch可以在华为的NPU芯片上运行
  

1.1 PyTorch的PRIVATEUSE1特性是什么？

• PyTorch不是一个完全的前后端分离的架构，但是可以从前后端分离的角度理解，其中前端负责用户交互，而后端处理计算任务；
• PyTorch的前端：主要有两个，一是Python前端：另一个是C++前端； Python前端是主要的接口，而C++前端为PyTorch机器学习框架提供了纯C++接口； Python API底层是基于C++代码库，提供了诸如张量和自动微分等基础数据结构和功能；
• PyTorch的后端：指的是执行前端定义操作的计算引擎。PyTorch支持多种后端，每种后端都针对特定的硬件或运行时进行了优化。例如，它有针对CUDA、cuDNN、MKL、MKLDNN、OpenMP等的后端。这些后端控制操作的行为，并且对于利用诸如GPU计算之类的硬件加速至关重要；
• PyTorch还允许自定义后端，用户可以使用torch.compile定义自己的后端。在追踪FX图后，可以通过TorchDynamo（torch.compile的图追踪组件）调用自定义后端中的后端函数；
• PRIVATEUSE1特性就是PyTorch提供的一种定制后端的机制，主要是为了帮助开发者在PyTorch中集成新的计算后端，PRIVATEUSE1特性为集成新的计算后端提供了一个结构化的方法。通过此特性，开发者可以将特定于硬件的优化和实现集成到PyTorch中，从而获得更好的性能，特别是在针对特定硬件加速器（如GPU、NPU或FPGA）时；
• 简单理解就是华为的Ascend NPU（昇腾NPU芯片）的OSS小组通过PyTorch的PRIVATEUSE1特性让PyTorch可以在华为的NPU芯片上运行。

1.2、 PyTorch2.1在华为NPU运行方式

参考官网，非常简单，代码如下：

torch.rename_privateuse1_backend("my_hardware_device")
torch.utils.generate_methods_for_privateuse1_backend()
x = torch.randn((2, 3), device='my_hardware_device')
y = x + x # run add kernel on 'my_hardware_device'

虽然通过插件（之前昇腾自己fork维护的）的形式官方支持了第三方硬件NPU，还有个极大的问题，我们可以看看PyTorch Adapter，版本配套关系及其离谱，就支持了三个pytorch版本，版本配套表一个电脑屏幕都放不下，昇腾在面向开发者文档上还是要多下功夫啊，感觉是研发根据自己的开发环境写的文档，和mindspore文档一样的问题（改天有时间再写），对于一个开发者文档是相当重要的，新手直接劝退…

另外，有意思的是，这个Pytorch Adapter在gitee上也开源了，但是文档不是同步更新，且内容逻辑是不一样的，大家可以去查查看；

gitee上的：PyTorch与Python版本配套表

github上的：版本配套表

2、 对pytorch和第三方芯片厂商（昇腾）有啥好处？

2.1 是原生支持了昇腾吗？

先通过【paperwithcode】看下目前市场上论文基于不同AI框架实现的框架占比情况来看，Pytorch依旧是作为一哥的存在；并且底层硬件基本都是英伟达，那么如果一个新硬件厂商想要支持Pytorch如何做呢？

我们先看看特斯拉咋做的：

正常来说：

1、 通过ONNX/torchscript作为中转来实现插件，但对于优化和开发效率有极大的影响；

2、 fork下pytorch，适配自己的硬件后并持续维护，但大家要知道Pytorch是按月发版本的，客户无法使用新特性，那硬件厂商的适配要吐血。。。

两种方式都会让使用新硬件的客户抓狂，那此次特性的更新，对于昇腾来说，其实是利好的，pytorch增加了PrivateUse1特性，相当于不用fork了！

但我们得看下其实也不是算原生支持，其实从1.1、1.2 的描述来看，pytorch增加了PrivateUse1特性，这样做可以降低对新硬 XPU 件的支持门槛，而PyTorch在做前后端做进一步的分离，以支持多硬件时代，芯片厂商实现后端后可以无缝切换；

为了验证这个功能呢，华为昇腾pytorch团队基于自己维护的torch_npu项目做了个新的后端，并且成功在torch 2.1中调用昇腾的npu;

注：彩蛋

GRAPHCORE发了一个PR（来源）增加IPU专用的DispatchKey通过了，有趣的是，一个多月后，昇腾希望在pytorch中加入NPU专用的DispatchKey（来源），但pytorch团队以PrivateUse1特性马上要支持了为由给拒绝了，昇腾还是慢了些啊，估计好气啊。。。

注：GRAPHCORE：拟未是一家人工智能芯片公司，为人工智能打造计算机系统，这些系统由先进的智能处理器（IPU）提供动力，旨在满足人工智能独特的计算要求。2016年，公司正式成立，总部位于英国布里斯托。2019年，公司在北京设立公司中国总部，并确定中文名“拟未”。目前，拟未已经设立了北京、上海、深圳和新竹办公室。

2.2对于pytorch的好处和挑战是什么？

• 好处：

进一步稳固pytorch的大哥的市场地位，难以撼动，pytorch本身也希望屏蔽硬件的差异，估计希望做到所有用户的一套代码对底层硬件无感知；估计后续pytorch的设备列表会越来越长…

• 挑战：

抛开对第三方硬件的支持这个特性，LLM领域对pytorch的冲击还是有的，当前pytorch对分布式训练的支持…

目前大厂都是pytorch+DeepSpeed/Megatron三件套在玩，大模型时代的对并行框架的需求已经是必需品了，如果只是一个pytorch那在LLM时代是没有意义的；

等等，哟，昇腾有AscendSpeed，据我所知MindSpore是把并行能力集成到框架本身，那么AscendSpeed就是DeepSpeed/Megatron等并行框架的插件咯，和适配pytorch是一个套路；好像也是没办法的办法~

2.3 对于第三方芯片厂商（昇腾）有的好处和挑战是什么？

• 好处：

对于pytorch来说是一小步，对于昇腾来说，解决了NPU支持pytorch的问题，算是生态上一大步，早就受不了第三方硬件上pytorch的各种适配和极低的开发效率；

• 挑战：

但torch_npu插件不是原生支持，会有几个较大的挑战：

1、版本更新：面对每季度release的pytorch，插件跟不上可能存在版本兼容性问题；

2、 第三方extension支持： 很多基于 torch 开发的项目本身也是一个 extension，默认一般都会有 CUDA 的支持，有的会有 CPU 的支持。也需要为新的硬件添加相应的kernel，可能的形式是做一个插件的插件，e.g. torchvision-xpu等等。如果前端 API 用法有些和 torch 不匹配的地方还要有自己的 modelzoo；

3、 **Test Coverage 的问题：**torch 本身的测试项目很多的，默认的测试不会跑非原生的 device，要保障质量是有一定难度的，毕竟 torch 用户那么多迭代了这么久才到现在的地步；

4、 LLM支持：大模型的爆发式增长，对新硬件的底层算子的支持有新的需求，如flashattension等大算子的支持，今天有flashattention，明天如果有另外一个算子的，昇腾还是要补齐基础算子的基础上，基于生态构建基于昇腾的优势算子，但昇腾的设备又没有2C端，终究在广大开发者生态上慢英伟达一拍。

加油啊，昇腾、加油啊dojo、加油啊~

天下苦英伟达久矣~

参考

1、 https://www.zhihu.com/question/624955377/answer/3239829901

2、 https://gitee.com/ascend/pytorch