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

3900 软件学报 2025 年第 36 卷第 8 期

用户对计算任务的需求. Hthreads 通过引入 DMA 技术实现高效的数据传输. 在优化案例中, Hthreads 性能显著优
于 MOCL3 [20] . 总体而言, Hthreads 的设计更贴近 MT-3000 芯片的底层架构, 在允许用户对 DSP 进行精细控制、
降低内核启动开销的同时, 也兼顾了可编程性和灵活性.
图 2 展示了 MT-3000 的各层级编程模型, 越接近底层编程模型效率越高, 但是其开发和维护难度也就越大.
综合考虑性能、可编程性、可读性等因素, 本文选择可以嵌入到 C 语言中的高级编程模型 Hthreads 作为开发
模型.

OpenCL
编
异构多线程库 (Hthreads) 程
可 模
编 高性能数学库 型
程 性
性 原语操作集 能
加速管理接口 (Iibmt) C+Intrinsic/汇编
汇编指令集
图 2 MT-3000 用户层编程模型

3 MTTorch 的实现与优化

3.1 MTTorch 设计背景
原生 PyTorch 框架中囊括了近千个算子, 为框架中所有算子都进行精心优化是一个庞大的工程. 考虑到大模型训
练对于计算需求的迫切性, 本文工作聚焦于在新一代天河超算上实现和优化 Transformer 类大语言模型训练任务.
首先, 我们在 MT-3000 的 CPU 上训练 PyTorch Examples (https://github.com/pytorch/examples) 官方项目中的
Transformer 模型至 loss 下降趋于稳定, 并使用 PyTorch 中的 profile 工具详细统计分析了模型训练过程中的各算
子的平均耗时 (如表 2 所示). 前 15 个算子的总耗时占比就高达 82% (如图 3 所示). 这为后续的算子优化提供了明
确的方向和依据. 本文将这 15 个算子定义为 Transformer 类模型的核心算子, 优化好这 15 个算子就可以极大地提
升模型的计算性能.

表 2 Transformer 类模型中耗时大的算子

函数名称 算子功能 时间占比 (%) 平均执行时间 (ms)
mm 矩阵乘法 46.20 19.454
softmax/log_softmax 归一化指数函数 12.47 141.763
add/sub/mul/div 加, 减, 乘, 除 11.24 3.87
softmax_back/log_softmax_back softmax的反向传播 3.95 44.959
bmm 批矩阵乘法 2.41 3.309
native layer_norm 层内归一化 2.00 2.147
sum Tensor求和 1.94 2.511
copy_ 数据拷贝 1.88 2.452
ReLU_back/GELU_back 激活函数反向传播 0.95 28.01
_cat Tensor拼接 0.82 3.23
ReLU/GELU 激活函数 0.57 15.91
表 2 中算子可分为 3 类. 第 1 类是如 cat、copy 等访存密集型操作, 并不涉及计算, 可以直接使用 DMA 进行
实现. 第 2 类为 pointwise 类型算子, 即 Tensor 中对应位置进行计算操作, 主要为计算耗时, 访存消耗较小. 第 3 类
为需要进行规约操作的算子, 如 layer_norm. 这一类算子在计算和访存方面都存在瓶颈问题, 且在规约过程中可能
遇到数值溢出. 综上所述, 本文面向 MT-3000 的体系结构特性, 对耗时占比高的 15 个核心算子进行了精心优化.