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

今日，寒武纪开源了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生态共建的决心。

一、 Torch-MLU的技术演进

1、设备后端接入PyTorch框架

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

▲侵入式设备后端支持

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

▲设备后端原生支持

通过这些改进，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开源生态注入新动力。同时通过技术的开源，赋予开发者更多能力，共同促进繁荣和多元的人工智能世界。