摘要:随着Transformer类大模型的飞速发展, 算力逐渐成为制约领域发展的瓶颈, 如何根据加速器硬件的结构特性加速和优化大语言模型的训练性能已成为研究热点. 面向天河新一代超算系统的加速芯片MT-3000, 提出并实现了适用于CPU+DSP异构架构的PyTorch扩展库——MTTorch, 其核心是一个多核并行的算子库, 对Transformer类模型训练过程中的核心算子进行向量化实现和优化. 同时, 针对MT-3000架构特性, 提出了面向多核 DSP 的高性能规约算法及乒乓算法, 显著提升了算子的运算性能. MTTorch还具有很好的通用性, 对于不同版本的 PyTorch都可以动态链接库的形式进行加载, 不改变PyTorch的原生实现. 大量实验证明, 实现的核心算子在 MT-3000 芯片上有着很好的性能, 在单DSP 簇上可以达到 8 倍的加速效果. 利用MTTorch在多节点执行训练任务时有着接近线性的加速比, 极大地提升了Transformer类模型在MT-3000 芯片上的训练效率.