IT之家 5 月 9 日消息,摩尔线程今日宣布推出 TorchMUSA v2.0.0 版本,这是其面向 PyTorch 深度学习框架的 MUSA 扩展库的重要升级。
新版本基于 MUSA Compute Capability 3.1 计算架构,支持原生 FP8 数据类型,支持 PyTorch 2.5.0,并通过多项针对 MUSA 计算平台的性能优化,进一步提升了对 AI 模型和大规模数据处理的支持能力。
作为本次升级的核心亮点,TorchMUSA v2.0.0 率先在国产 GPU 上实现了对 FP8 数据类型的完整支持。
FP8 是当前 AI 计算的一种前沿低精度格式,在支持原生 FP8 的 GPU 上,大语言模型(LLM)训练采用 FP8 混合精度可大幅提高 GPU 算力,显著降低显存占用。
摩尔线程基于新一代 MUSA Compute Capability 3.1 计算架构的全功能 GPU 原生支持 FP8 计算,这为 TorchMUSA v2.0.0 实现 FP8 矩阵乘法和分布式通信优化提供了基础。依托这一底层架构优势,TorchMUSA v2.0.0 能够充分发挥 FP8 的计算效能,提升大语言模型训练和推理的效率。
TorchMUSA v2.0.0 在 MUSA 计算平台引入多项创新功能,进一步提升深度学习任务的执行效率:
新增虚拟内存管理支持:
MUSA 虚拟内存管理技术能够有效缓解 GPU 内存碎片化问题,降低模型训练过程中的峰值内存占用,特别适用于 FSDP、DeepSpeed 和 MegatronLM 等主流大模型训练框架。
新增 MUSA Graph 支持:
MUSA Graph 技术将多个 MUSA 内核整合到一个图中,通过单次 CPU 调度大幅减少启动开销,提升计算效率,同时与 CUDA Graph 接口高效兼容。
torch.compile 增加 Triton 后端支持:
为 torch.compile 提供了 TritonMUSA 后端支持,开发者可以直接使用 PyTorch 原生接口,获得更高效的性能表现。
TorchMUSA v2.0.0 在完整支持 PyTorch 2.2.0 的基础上,新增了对 PyTorch 2.5.0 的支持,使开发者能够在基于 MUSA Compute Capability 3.1 计算架构的全功能 GPU 上,无缝运行新版本的 PyTorch。
TorchMUSA 已完全开源,开发者可通过访问 GitHub 获取源代码。IT之家附开源地址: