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

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

RQ1: MTTorch 中的算子的正确性如何?
RQ2: MTTorch 中的虚拟算子是否避免了引入额外开销?
RQ3: MTTorch 中核心算子的性能提升效果如何?
RQ4: 使用 MTTorch 进行模型训练时的性能和加速比如何?
本文在由 MT-3000 芯片构建的天河新一代超算系统上进行测试. 使用 PyTorch 1.10 (https://pytorch.org/
docs/1.10/) 版本作为软件环境, 以 CPU 同构环境下的 PyTorch 作为基准线, 并在编译过程中集成了 OpenBLAS
(https://github.com/xianyi/OpenBLAS) 等 CPU 上的高性能库. 为公平起见, 本文在单个 MT-3000 片上用 4 个 CPU
核和单个 DSP 簇进行比较. 区别于科学计算, 深度学习训练任务的一大不同在于任务对加速器内存的需求更大.
MT-3000 片上 DDR 总量为 64 GB, 常规状态下, 每块 DSP 可分得的内存大小为 12 GB 左右. 为了保证大语言模型
训练可以更稳定地运行, 只启用一块 DSP 簇, 保证其可支配 48 GB 内存, 防止内存溢出 (out of memory, OOM). 在
算子测试阶段, 采用 Python 的单元测试框架 unittes (https://docs.python.org/3/library/unittest.html) 进行单元测试,
并使用 PyTorch 中的 profile 工具对算子性能进行分析. 在模型测试阶段, 本文在单节点和多节点情况下分别进行
了模型训练性能和可扩展性测试. 考虑到 MTTorch 的设计模式, 本文的方法也适用于其他 PyTorch 版本.

4.1 RQ1: MTTorch 中算子的正确性
unittest 单元测试提供了创建测试用例, 测试套件以及 批量执行的方案, 用户安装 Python 之后可以通过
import unittest 引入并直接使用. 利用 unittset 包, 对 15 个核心算子和 21 个大模型所需的虚拟算子在 3 个不同数据
规模下进行了单元测试, 基准为 CPU 算子的计算结果. 本文采用 torch.testing.assert_close 判断结果是否相等, 相对
误差阈值设置为 1E–06, 这是主流深度学习框架进行测试时候采用的通用值. 由于虚拟算子的实现逻辑一致, 因此
在实验中以 nll_loss_forward 为代表进行说明, *表示该算子为虚拟算子. 表 3 展示了 MTTorch 中的算子在不同数
据规模下的正确性测试结果, 其中√表示算子通过了正确性测试, 实验结果为 5 次随机实验后综合得出.

表 3 MTTorch 算子正确性测试

测试算子 2 18 2 24 2 28
add/sub √ √ √
mul/div √ √ √
GELU √ √ √
ReLU √ √ √
GELU_back √ √ √
ReLU_back √ √ √
softmax √ √ √
log_softmax √ √ √
softmax_back √ √ √
log_softmax_back √ √ √
layer_norm √ √ √
mm √ √ √
sum √ √ √
nll_loss_forward * √ √ √

由结果可得, 无论是核心算子还是虚拟算子, 计算误差都小于 1E–06.

4.2 RQ2: MTTorch 中虚拟算子的性能
22
在本节中, 对比了虚拟算子和 CPU 同构环境下 PyTorch 算子的性能. 数据规模为 2 个单精度浮点数.
从图 9 可以看出, 虚拟算子性能和 CPU 实现的算子不相上下, 这表明通过修改设备属性改变算子运算后端的
方法几乎未引入额外开销, 达到了我们的设计目标. 同时, 在大多数情况下, 虚拟算子的性能还好于 CPU 算子. 这