Page 490 - 《软件学报》2025年第8期

王昊天 等: MTTorch: 面向 MT-3000 芯片和 Transformer 模型的 PyTorch 算子库实现与优化 3913

MT CPU 500 MT CPU
700 DSP 向量化 DSP 向量化
DSP 向量化+乒乓算法 DSP 向量化+乒乓算法
600 400
执行时间 (ms) 400 执行时间 (ms) 300
500

300
200
200
100
100
0 0
2 18 2 20 2 22 2 24 2 26 2 28 2 18 2 20 2 22 2 24 2 26 2 28
数据量 数据量
(a) log_softmax (b) layer_norm

MT CPU 400 MT CPU
350
DSP 向量化 DSP 向量化
DSP 向量化+乒乓算法 350 DSP 向量化+乒乓算法
300
300
执行时间 (ms) 200 执行时间 (ms) 250
250
200
150
100 150
100
50 50
0 0
2 18 2 20 2 22 2 24 2 26 2 28 2 18 2 20 2 22 2 24 2 26 2 28
数据量 数据量
(c) softmax_back (d) softmax

500 MT CPU
DSP 向量化
DSP 向量化+乒乓算法
400
执行时间 (ms) 300

200

100

0
2 18 2 20 2 22 2 24 2 26 2 28
数据量
(e) log_softmax_back
图 12 算子优化消融实验

可以看出, 两个模型的加速比都要显著优于 CPU 同构环境下的 PyTorch, 且随着节点的增加加速比逐渐呈线
性增加, 说明本文实现的方法可以为基于 MT-3000 搭建的高性能计算集群带来更强的可扩展性, 使得 DSP+CPU
异构集群在人工智能领域的潜力可以被进一步激发.

5 总 结

本文结合训练大语言模型的实际需求, 介绍了面向天河新一代超算系统加速芯片 MT-3000 的 PyTorch 扩展
库——MTTorch 的设计和实现. MTTorch 运用单例模式对 DSP 设备和线程组进行高效管理. 算子层面, 本文利用