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

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

引用

CSDN

1.

https://blog.csdn.net/sinat_36458870/article/details/146039311

PyTorch 2.1版本正式宣布支持华为昇腾芯片，这一技术突破不仅为昇腾芯片的生态建设注入了新的动力，也体现了PyTorch在多硬件支持方面的持续努力。本文将深入探讨这一变化的技术细节及其对PyTorch和昇腾芯片的影响。

1. 发生了什么事儿？

在2023年10月4日PyTorch 2.1版本的发布博客上，PyTorch介绍了一个重要的beta版本新特性：提高了第三方设备的支持。具体来说，华为Ascend NPU（昇腾NPU芯片）的OSS小组已经成功将torch_npu整合进入PyTorch。这意味着PyTorch原生支持的硬件范围进一步扩大，昇腾NPU现在也可以直接运行PyTorch代码。

1.1 PyTorch的PRIVATEUSE1特性是什么？

PyTorch的架构虽然不是完全的前后端分离，但可以从这个角度来理解。前端负责用户交互，而后端处理计算任务。

• 前端：主要有两个，一是Python前端，另一个是C++前端。Python前端是主要的接口，而C++前端为PyTorch机器学习框架提供了纯C++接口。Python API底层是基于C++代码库，提供了诸如张量和自动微分等基础数据结构和功能。

• 后端：指的是执行前端定义操作的计算引擎。PyTorch支持多种后端，每种后端都针对特定的硬件或运行时进行了优化。例如，它有针对CUDA、cuDNN、MKL、MKLDNN、OpenMP等的后端。这些后端控制操作的行为，并且对于利用诸如GPU计算之类的硬件加速至关重要。

• 自定义后端：PyTorch允许用户通过torch.compile定义自己的后端。在追踪FX图后，可以通过TorchDynamo（torch.compile的图追踪组件）调用自定义后端中的后端函数。

• PRIVATEUSE1特性：这是PyTorch提供的一种定制后端的机制，主要是为了帮助开发者在PyTorch中集成新的计算后端。PRIVATEUSE1特性为集成新的计算后端提供了一个结构化的方法。通过此特性，开发者可以将特定于硬件的优化和实现集成到PyTorch中，从而获得更好的性能，特别是在针对特定硬件加速器（如GPU、NPU或FPGA）时。

1.2 在华为NPU上运行PyTorch

参考官方文档，使用昇腾NPU运行PyTorch代码非常简单，以下是示例代码：

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版本。此外，昇腾在面向开发者文档上还需要进一步优化，新手可能会因此感到困扰。

另外，值得注意的是，这个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