软件学报

作者：王昊天，孙羽菲，隋轶丞，王嘉豪，石昌青，方建滨，张玉志

单位：王昊天，南开大学 软件学院, 天津 30045011,孙羽菲，南开大学 软件学院, 天津 30045002,隋轶丞，南开大学 软件学院, 天津 30045003,王嘉豪，南开大学 软件学院, 天津 30045004,石昌青，南开大学 软件学院, 天津 30045005,方建滨，国防科技大学 计算机学院, 湖南 长沙 41007306,张玉志，南开大学 软件学院, 天津 30045007

关键词：PyTorch;高性能计算;Transformer模型;天河超级计算机;CPU+DSP异构计算;软件生态

基金：国家重点研发计划(2021YFB0300104); 先进计算与关键软件海河实验室科技项目(22HHXCJC00001); 启元实验室创新基金(2022-JCJO-LA-001-068)

随着Transformer类大模型的飞速发展, 算力逐渐成为制约领域发展的瓶颈, 如何根据加速器硬件的结构特性加速和优化大语言模型的训练性能已成为研究热点. 面向天河新一代超算系统的加速芯片MT-3000, 提出并实现了适用于CPU+DSP异构架构的PyTorch扩展库——MTTorch, 其核心是一个多核并行的算子库, 对Transformer类模型训练过程中的核心算子进行向量化实现和优化. 同时, 针对MT-3000架构特性, 提出了面向多核 DSP 的高性能规约算法及乒乓算法, 显著提升了算子的运算性能. MTTorch还具有很好的通用性, 对于不同版本的 PyTorch都可以动态链接库的形式进行加载, 不改变PyTorch的原生实现. 大量实验证明, 实现的核心算子在 MT-3000 芯片上有着很好的性能, 在单DSP 簇上可以达到 8 倍的加速效果. 利用MTTorch在多节点执行训练任务时有着接近线性的加速比, 极大地提升了Transformer类模型在MT-3000 芯片上的训练效率.

来源：2025年第8期

《软件学报》期刊编辑部

查看软件学报杂志2025年第8期