今日,寒武纪开源了PyTorch设备后端扩展插件Torch-MLU,并实现了寒武纪硬件对于PyTorch的原生支持,充分提升了开发者的使用体验和集成效率。

早在2018年,寒武纪就开始了与PyTorch框架的集成适配工作。自PyTorch 1.3版本开始就支持寒武纪MLU系列智能加速卡作为PyTorch的加速后端,使开发者能够将原先基于GPU的深度学习网络高效迁移到基于寒武纪MLU的PyTorch环境中。

随着PyTorch社区发布PyTorch 2.4版本,配合新版本(≥ V1.22)的Torch-MLU插件,寒武纪硬件实现了对PyTorch的原生支持。开发者无需再安装寒武纪定制的Cambricon PyTorch,而是可以基于社区的PyTorch安装Torch-MLU插件,就可将寒武纪硬件的算力接入原生PyTorch。

与此同时,Torch-MLU的开源使寒武纪能够更快地响应开发者的反馈和问题,搭建了高效的沟通桥梁,促进了与全球开发者的协作和知识共享。寒武纪通过持续贡献开源社区的实际行动,展示着拥抱开源、支持AI生态共建的决心。

GitEE仓库(点击文末“阅读原文”可直接跳转)↓

https://gitee.com/cambricon/torch_mlu

*具体版本与分支信息请参考仓库中的README文档。

一、 Torch-MLU的技术演进

1、设备后端接入PyTorch框架

PyTorch初期没有提供第三方后端统一接入方案。新设备后端如果想接入PyTorch,往往需要对PyTorch的核心代码进行修改。比如开发者需要在核心模块Dispatcher中加入对应后端的Dispatch Key,以及对其他模块中关于设备后端做一些适配性修改。设备后端扩展插件需要配合修改过的PyTorch才能够正常工作。

原生PyTorch支持,大模型一键迁移!寒武纪开源Torch-MLU

▲侵入式设备后端支持

2023年8月PyTorch社区发布了PyTorch 2.1版本,重磅推出了基于PrivateUse1 Dispatch Key的新后端统一接入参考设计。PrivateUse1设计方案通过PrivateUse1抽象层将设备厂商私有属性统一封装,完善了PyTorch的设备扩展机制,为新硬件非侵入性接入PyTorch提供了途径。

原生PyTorch支持,大模型一键迁移!寒武纪开源Torch-MLU

▲设备后端原生支持

通过这些改进,PyTorch不仅增强了其作为深度学习框架的灵活性和可扩展性,还为开发者和设备厂商提供了更顺畅和高效的合作环境。

2、寒武纪持续贡献PyTorch社区

PrivateUse1接入方案的推出为新后端统一接入提供了启发性的思路,但是在初期仍有一些待完善的特性和模块,需要设备厂商的持续投入和贡献。

PrivateUse1支持设备在Eager模式下的基本的运行和接入,但并不能完善支持PyTorch的一些特性。比如Profiler模块,初期仅支持GPU后端,不支持外部自定义后端,基于PrivateUse1注册的设备后端无法使用PyTorch原生的Profiler。

PrivateUse1路径的限制给开发者带来了更多挑战,使他们在进行设备集成时面临更多障碍,影响了整体使用体验。

寒武纪在PrivateUse1接入方案的基础上优化了非CUDA设备的接入体验,在今年向PyTorch社区提交了数十个Patch,涉及Profiler、Compile、Graph Capture、Autograd、Allocator、Storage, FSDP、Sparse等众多模块,打通了这些模块与PrivateUse1的集成路径,进一步完善了PrivateUse1机制。未来,第三方设备的厂商和开发者就能够充分利用这些已打通的集成路径,更加高效和便捷地接入PyTorch。

3、安装Torch-MLU插件,保留原生体验

经过寒武纪对PyTorch社区的PrivateUse1机制进行完善后,自PyTorch 2.4版本起,寒武纪硬件实现了PyTorch原生支持。

开发者可直接基于社区PyTorch发布包安装Torch-MLU插件,即可将寒武纪MLU的算力无缝接入,同时保留原生PyTorch的使用体验。

二、开源范围

当前开源的Torch-MLU仓库支持的社区版本:

  • PyTorch 2.1
  • PyTorch 2.3
  • PyTorch 2.4(2024年7月25日最新发布)

未来会随社区版本持续更新。

三、拥抱开源社区,持续生态贡献

寒武纪长期以来秉承开放、合作、共享的理念,积极参与开源社区的建设,在多个重要开源项目中贡献代码,如PyTorch、TensorFlow、Huggingface、Transformers、vLLM、Deepspeed等大模型训练推理应用中的核心组件。

近期,寒武纪开源了Triton-Linalg AI编译器前端,开发者或者硬件厂商可以以极低的开发成本,快速集成支持Triton语言特性的后端指令集,并对接上层AI应用。

此次开源Torch-MLU插件,也是希望未来能更好地理解和更快速地解决开发者的问题,同时为寒武纪深度学习框架与开发者之间建立直接的交流渠道。寒武纪坚信,推动人工智能领域未来发展的关键在于开放性、合作精神和共享资源。通过开源项目,开发者们可以团结一致,通过协作和配合,使得技术成果得到广泛应用。

未来,寒武纪将持续贡献开源社区,为AI开源生态注入新动力。同时通过技术的开源,赋予开发者更多能力,共同促进繁荣和多元的人工智能世界。