自有AI芯片接入AI框架Pytorch的方案

现在AI框架主要用Pytorch,包括一些常用的库对Pytorch支持都较好,国产AI芯片很多也对AI框架Pytorch进行了支持

一、华为昇腾npu

能够跟上Pytorch的更新,直接和Pytorch兼容,而且有较多人来维护,代码风格不错,之前是通过注入Pytorch预留的的xla搞的接入,Pytorch 2.1起被官方接收了。这个代码非常推荐,笔者用他的框架实现了自有GPGPU芯片手写数字识别的训练,2023.8跑路了,靠只接入算子是能跑的,但是不做fuse效率很差。

二、寒武纪mlu

这个版本很老,他们是通过docker为客户提供运行环境的,pytorch_patches文件夹表明他们需要修改Pytorch源代码

三、谷歌tpu

Pytorch 给他留了接口,甚至aten中也有为了支持他的某些feature的源码,tpu也算DSA的引领者

四、OpenCL 实现

一个老哥在孜孜不倦得推动OpenCL在Pytorch的适配,他还一直在维护,可惜OpenCL都要无了(

五、摩尔线程 MTT

2024.1更新 当时团队找错了,正确的是MooreThreads,发布会的时候应该还没公开,2023年6月底放上来的,接入方式和寒武纪差不多的
不知道他开源在哪里了,github团队是空的,可能合作商能从他官网拿到?他说自己开源了, 2023夏季发布会

六、Inductor backend


Pytorch2.0 提供的新的接入方式,官方文档,Triton就是这样的,你就可以省掉些算子开发,不过从MLIR还要再接你的编译器后端,但是他能快速跑完流程。nvFuser就被Triton替代了,Triton很稳。

七、独树一帜支持国产框架 PaddlePaddle

PaddlePaddle还是有不少国产芯片提供了支持的,不过要谨防PaddleNLPPaddle版本间的不匹配问题。但是『飞桨』自定义硬件接入实现 还是比较nice的

八、其他

沐曦是有的,看到这个pr里提到了,应该也是给的docker
天数智芯iluvatar 也有自己维护的库,不过找不到他们的源码
昆仑芯 有适配,疑似源码,只是克隆了过来
壁仞Github团队 团队有的,不过没有pytorch源码
燧原Github团队 团队有的,没有pytorch源码,但是他们用EnflameGCU对飞桨进行了支持。torch_gcu 已经有了,疑似源码

posted @ 2023-12-06 12:31  暴力都不会的蒟蒻  阅读(1167)  评论(0编辑  收藏  举报