今日，寒武纪开源了 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 才能够正常工作。

▲侵入式设备后端支持

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