华为大模型生态重要一步！PyTorch最新2.1版本宣布支持华为昇腾芯片（HUAWEI Ascend）

大模型的发展一个重要的基础条件是底层硬件计算能力的大幅提高，特别是GPU的发展，与transformer架构的大模型训练非常契合。当前全球最大的GPU供应商英伟达系列的显卡几乎垄断了大模型训练与推理的所有GPU芯片市场。除了英伟达显卡本身算力强悍外，基于英伟达GPU之上构建的CUDA、PyTorch等平台软件生态也是非常重要的一环。而最新的PyTorch2.1版本发布的一个beta特性中包含了对华为昇腾芯片的支持，这也是大模型生态多样性发展的一个很重要的信号。

PyTorch 2.1对华为昇腾芯片NPU的支持

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

PyTorch的PRIVATEUSE1特性简介

PRIVATEUSE1特性是PyTorch中的引入新特新，主要目标是让PyTorch可以使用更多的硬件。体现了该框架在适应不断变化的硬件和计算环境方面的努力。

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芯片上运行。

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'