Page 475 - 《软件学报》2025年第8期
P. 475
3898 软件学报 2025 年第 36 卷第 8 期
本文第 1 节给出深度学习算子和框架在不同设备上转换适配的相关工作. 第 2 节主要介绍了 MT-3000 体系
架构以及编程模型的选择. 第 3 节介绍了 MTTorch 的整体架构以及核心算子的异构并行编程实现. 第 4 节设计实
验验证了 MTTorch 在 MT-3000 芯片上的正确性和加速效果. 第 5 节对本文的工作进行总结并提出了未来的研究
方向.
1 机器学习算子及框架在异构设备上扩展的相关工作
近年来, 大语言模型技术得到了飞速发展, 而语言模型的研究和应用却可能遭遇算力瓶颈. 传统应用的算力支
撑更多地依靠 CPU, 而 AI 所需的智能算力在 CPU 之外, 还需要大量 GPU 参与. 为充分发挥通用加速器, 尤其是
国产高性能异构计算集群在人工智能领域的巨大潜力, 众多研究者在算子和框架的移植优化领域做了大量工作.
李亿渊等人 [13] 通过自定义的矩阵分层分块技术, 在申威 26010 异构众核处理器上实现了高效的稀疏矩阵乘
法操作. Yin 等人 [14] 在新一代数字信号处理器 FT-M7032 上结合硬件特性, 实现了一种可以自动装配微内核, 自动
调整分块大小和并行策略的非规则通用矩阵乘法 (GEMM). 裴向东等人 [15] 针对飞腾异构多核 DSP 的体系架构特
点与矩阵转置操作的特点, 利用 DSP 模 16 存储功能和隐式乒乓设计实现了适用于多种位宽的异构并行矩阵转置
算法 ftmMT. 王庆林等人 [16] 针对飞腾多核处理器对 Winograd 快速卷积算法进行优化. 上述方法都是针对深度学
习或科学计算中的某一个具体的算子, 如何在基于 DSP 的国产高性能异构设备上高效地使用深度学习框架训练
模型还有待解决.
随着深度学习领域对算力需求的日益增长, 越来越多的计算设备被运用于人工智能领域. 如何将基于 CUDA
的深度学习框架移植到不同的设备上受到了大量关注. OpenCL Caffe 是 AMD 开发的基于 OpenCL 的并行深度学
习框架 [17] , 它可以利用不同平台上的硬件资源来实现高效的计算. 它具有跨平台、高效能的优点, 可以在不同的
硬件设备上运行, 例如 GPU、CPU 和 FPGA 等. OpenCL Caffe 还提供了一些优化技术, 例如基于局部内存的优化、
图像数据的预处理和统一的内存访问等, 以提高性能和效率. Chen 等人 [18] 针对 MT-3000 处理器提出了一个高效
的卷积神经网络推理框架 PCIF, 该框架支持通过 ONNX 导入和优化训练好的深度学习模型, 通过计算图融合和
多核并行向量化技术提高计算效率. 实验结果显示, PCIF 能有效加速神经网络的推理过程. 然而, 该框架目前仅实
现了卷积神经网络中常用算子的前向传播功能, 尚未实现相关算子的反向传播功能. 因此 PCIF 仅适用于推理场
景, 不支持模型训练. 此外, 该框架只适配了卷积网络相关的算子, 不支持 Transformer 类大语言模型的训练. 陈锐
等人 [19] 探索了 CUDA 核函数向 OpenCL 核函数转换的基本流程, 并通过单例模式初始化, 缓存编辑文件以及使用
规约操作代替原子操作等优化方式, 实现了基于 OpenCL 的 TensorFlow, 在 GPU、华为鲲鹏-920CPU 等设备上具
有通用性. 然而, 由于 OpenCL 的通用性, 算子的实现难以针对特定的体系架构进行优化, 且当前对于 PyTorch 框
架如何移植到 CPU+DSP 异构环境上还有待研究, 导致现有工作均无法充分发挥 DSP 异构芯片的性能.
本文的工作面向 MT-3000 处理器的异构众核架构, 提供了支持 DSP+CPU 异构训练的 PyTorch 扩展库——
MTTorch, 针对训练 Transformer 类模型常用的 15 个算子, 设计了并行算法. 大量实验证明了 MTTorch 相比基于
CPU 同构环境的 PyTorch 有着更好的性能和可扩展性.
2 MT-3000 软硬件架构
2.1 MT-3000 硬件体系架构
MT-3000 是一款高性能异构多区域处理器, 由多核 CPU 和 4 个 DSP 簇构成, 是天河新一代超算系统的重要
加速部件. 其 CPU 类似于简化版的 FT-2000 Plus 处理器, 具有 16 个兼容 ARMv8 架构的 CPU 核心, 用于进程管理
和通信, 并运行操作系统. MT-3000 的加速核被平均分为 4 个自动加速集群 (DSP), 每一个 DSP 簇包含 24 个由 1
个控制核心和 16 个加速器核心组成的 DSP 核、一个大小为 6 MB 的全局共享内存 (global shared memory, GSM)、
高带宽共享内存 (HBSM) 和芯片上的 DDR 内存空间. MT-3000 体系架构如图 1 所示.

