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

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

pong algorithm for multi-core DSP, significantly improving the computational performance of the operators. MTTorch also has good
generality as it can be loaded as a dynamic link library for different versions of PyTorch without changing the native implementation of
PyTorch. Extensive experiments show that the core operators implemented in this study have excellent performance on MT-3000 chip,
achieving 8 times acceleration on a single DSP cluster. Using MTTorch for training tasks on multiple nodes achieves nearly linear
acceleration, greatly improving the training efficiency of Transformer-based models on MT-3000 chip.
Key words: PyTorch; high performance computing; Transformer model; Tianhe supercomputer; CPU+DSP heterogeneous computing; software
ecology

随着深度学习技术的发展, 深度神经网络取得了长足的进步, 模型的参数量以及训练数据也在飞速增长. 2020
[1]
年, OpenAI 发布的全球最大的预训练语言模型——GPT-3 参数量甚至达到了惊人的 1 750 亿, 2022 年发布的
[2]
ChatGPT 更是把大模型提升到了全新的高度. 当前的大语言模型都是在 Transformer 模型 [3,4] 的基础上改进而来.
然而, Transformer 模型的大规模训练需要巨大的计算资源. 因此, 如何合理利用已有计算设备, 尤其是国产大规模
计算集群, 提高计算效率成为了一个关键问题 [5,6] .
为了提高深度学习模型对于计算设备的利用率以及简化模型搭建流程, 让研究人员更专注于模型设计, 众多
知名互联网公司设计了多种机器学习框架. PyTorch 是一个基于 Python 的机器学习框架, 由 Facebook 开发并维
护, 其采用命令式编程, 可见及所得, 接口相对比较统一, 使得用户能够更加灵活地构建和训练深度学习模型. 同
时, PyTorch 还提供了丰富的工具包和 API, 使得用户能够轻松地进行数据处理、性能分析和调试等操作, 因而在
工业界和学术界都得到了广泛的运用. 然而, 当前 PyTorch 官方只支持 CUDA (compute unified device arch-
itecture) 和 [7] ROCm (radeon open compute) (https://rocmdocs.amd.com/en/latest/), 故而只能使用英伟达和 AMD 生产
的 GPGPU 进行计算, 导致众多异构计算集群无法高效运行 PyTorch 任务, 令其在人工智能领域的巨大潜力难以
发挥.
出于能源和功耗方面的考虑, 最先进的高性能计算平台往往采用低功耗嵌入式架构, 如数字信号处理器
[8]
(digital signal processor, DSP). MT-3000 是国防科技大学研发的一款用于高性能计算的异构多区域处理器, 搭载
了 4 个 DSP 簇. 与 CPU、GPU 和 TPU 等设备不同, MT-3000 采用超长指令集 (very long instruction word, VLIW)
和加速阵列微架构, 每个加速器核心以 VLIW 的方式工作, 每 16 个加速器核心和一个控制核心被组织成一个加速
[9]
阵列. 同时其也集成了直接存储器取存引擎 (direct memory access, DMA) 用于内存操作. 基于 MT-3000 搭建的
天河新一代超算系统 [10] 在高性能计算和深度学习领域都有着巨大的潜能. 同时, MT-3000 具备了完整的编程工具
链和向量 C 语言支持, 使得用户可以在 C 语言层面充分利用 VLIW 和加速阵列微体系架构的组合潜力. 研究表
明, 只有软件能够充分利用硬件特性, 加速器才能释放其潜力 [11,12] . 因此, 面向新一代 DSP 加速设备开发 PyTorch
支持是利用众多高性能计算平台进行模型训练的必要一步.
本文面向 MT-3000 芯片的体系架构特点, 实现了一个 PyTorch 的扩展库——MTTorch. 由于 PyTorch 中包含
的算子数量庞大, 结合训练大语言模型的需求, 本文聚焦于 Transformer 类模型, 对模型训练过程中的 15 个核心算
子在 CPU+DSP 异构环境下进行精心优化.
本文的创新点和贡献如下.
(1) 提出并实现了面向 MT-3000 异构加速器的 PyTorch 扩展库——MTTorch. 结合 MT-3000 的体系结构特
点, MTTorch 对 Transformer 类模型训练中的核心算子进行了实现和并行优化, 为新一代天河超算支持大模型训练
打下了坚实的基础.
(2) 针对 MT-3000 体系结构的特点, 提出了面向多核 DSP 的高性能规约算法; 利用核内 DMA 和计算部件异
步执行特性设计了乒乓算法. 基于上述优化策略, 并结合 DSP 上的向量寄存器特性, 有效提升了计算效率.
(3) 设计实现了面向 DSP 的设备管理模块以及线程组管理模块, 基于单例模式有效避免了 DSP 簇重复启闭造
成的额外开销.
(4) 通过一系列实验, 验证了 MTTorch 所优化实现的算子均能正确运行, 并且获得了 2–8 倍的性能提升. 同时,
基于 MTTorch 的大模型训练任务展现出了接近线性的加速比.