寒武纪开源Torch-MLU：实现原生PyTorch支持与大模型一键迁移，助力AI生态发展

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

早在2018年，寒武纪便开始了与PyTorch框架的集成适配工作。从PyTorch 1.3版本开始，寒武纪MLU系列智能加速卡就被纳入PyTorch的加速后端，使得开发者能够高效地将原本基于GPU的深度学习网络迁移到寒武纪MLU平台。

随着PyTorch 2.4版本的发布，结合新版本（≥ V1.22）的Torch-MLU插件，寒武纪硬件实现了对PyTorch的原生支持。现在，开发者不再需要安装定制版的Cambricon PyTorch，只需基于官方发布的PyTorch安装Torch-MLU插件，即可轻松接入寒武纪硬件的算力。

Torch-MLU的开源不仅加快了寒武纪对开发者反馈和问题的响应速度，还搭建了高效的沟通桥梁，促进了与全球开发者的协作和知识共享。寒武纪通过实际行动展示了其拥抱开源、支持AI生态共建的决心。

GitEE仓库链接：
https：//gitee.com/cambricon/torch_mlu
*具体版本与分支信息请参阅仓库中的README文档。*

Torch-MLU的技术演进

1. 设备后端接入PyTorch框架

PyTorch早期并未提供第三方后端统一接入方案，新设备后端通常需要对PyTorch的核心代码进行修改。例如，在核心模块Dispatcher中加入对应的Dispatch Key，并对其他模块进行适配性修改。这种设备后端扩展插件需要配合修改过的PyTorch才能正常工作。

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

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

PrivateUse1接入方案虽然为新后端统一接入提供了启发性的思路，但在初期仍存在一些待完善的特性和模块。例如，Profiler模块最初仅支持GPU后端，不支持外部自定义后端。这给开发者带来了很多挑战，影响了整体使用体验。

寒武纪在此基础上优化了非CUDA设备的接入体验，向PyTorch社区提交了数十个Patch，涉及Profiler、Compile、Graph Capture、Autograd、Allocator、Storage, FSDP、Sparse等多个模块，进一步完善了PrivateUse1机制。这些改进为第三方设备厂商和开发者提供了更加高效便捷的接入方式。

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

本文来源： 智东西【阅读原文】