快科技 5 月 9 日消息，近日，国产 GPU 厂商摩尔线程正式发布 Torch-MUSA v2.0.0 版本，这是其面向 PyTorch 深度学习框架的 MUSA 扩展库的重要升级。

在 Torch-MUSA 中，用户只需指定 torch.device ( "musa" ) ，即可将现有的 PyTorch 模型迁移到 MUSA 架构的 GPU 上运行，无需大幅修改代码，目前 Torch-MUSA 已完全开源，可通过 GitHub 获取源代码。

作为本次升级的核心亮点，Torch-MUSA v2.0.0 率先在国产 GPU 上实现了对 FP8 数据类型的完整支持。

FP8 是当前 AI 计算的一种低精度格式，在支持原生 FP8 的 GPU 上，大语言模型训练采用 FP8 混合精度可大幅提高 GPU 算力，降低显存占用。

摩尔线程基于新一代 MUSA Compute Capability 3.1 计算架构的全功能 GPU 原生支持 FP8 计算，为 Torch-MUSA v2.0.0 实现 FP8 矩阵乘法和分布式通信优化提供了基础。

依托这一底层架构优势，Torch-MUSA v2.0.0 能够充分发挥 FP8 的计算效能，显著提升大语言模型训练和推理的效率。

Torch-MUSA v2.0.0 在 MUSA 计算平台引入多项创新功能，进一步提升深度学习任务的执行效率，主要包括：

1、新增虚拟内存管理支持：

MUSA 虚拟内存管理技术能够有效缓解 GPU 内存碎片化问题，降低模型训练过程中的峰值内存占用，特别适用于 FSDP、DeepSpeed 和 Megatron-LM 等主流大模型训练框架。

2、新增 MUSA Graph 支持：

MUSA Graph 技术将多个 MUSA 内核整合到一个图中，通过单次 CPU 调度大幅减少启动开销，提升计算效率，同时与 CUDA Graph 接口高效兼容。

3、torch.compile 增加 Triton 后端支持：

为 torch.compile 提供了 Triton-MUSA 后端支持，开发者可以直接使用 PyTorch 原生接口，获得更高效的性能表现。

不仅如此，Torch-MUSA v2.0.0 在完整支持 PyTorch 2.2.0 的基础上，还新增了对 PyTorch 2.5.0 的支持，使开发者能够在基于 MUSA Compute Capability 3.1 计算架构的全功能 GPU 上，无缝运行新版本的 PyTorch。

未来 Torch-MUSA 还将继续跟进 PyTorch 的版本更新，计划支持更高版本的 PyTorch。