MTTorch: PyTorch Arithmetic Library Implementation and Optimization for MT-3000 Chip and Transformer Model
Author:
Affiliation:

Clc Number:

Fund Project:

  • Article
  • |
  • Figures
  • |
  • Metrics
  • |
  • Reference
  • |
  • Related
  • |
  • Cited by
  • |
  • Materials
  • |
  • Comments
    Abstract:

    With the rapid development of Transformer-based large models, computing power has gradually become a bottleneck in the development of this field. Research hotspots rely on how to accelerate and optimize the training performance of large language models based on the structural characteristics of accelerator hardware. This study proposes and implements MTTorch, a PyTorch extension library for the CPU+DSP heterogeneous architecture, which is applicable to the MT-3000 accelerator chip of the new generation of the Tianhe supercomputer. The core of MTTorch is a multi-core parallel operator library that vectorizes and optimizes the core operators during the training of Transformer-based models. Additionally, this study innovatively proposes a high-performance reduction algorithm and a ping-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.

    Reference
    Related
    Cited by
Get Citation

王昊天,孙羽菲,隋轶丞,王嘉豪,石昌青,方建滨,张玉志. MTTorch: 面向MT-3000芯片和Transformer模型的PyTorch算子库实现与优化.软件学报,2025,36(8):3896-3916

Copy
Share
Article Metrics
  • Abstract:
  • PDF:
  • HTML:
  • Cited by:
History
  • Received:November 13,2023
  • Revised:March 13,2024
  • Adopted:
  • Online: December 31,2024
  • Published: August 06,2025
You are the firstVisitors
Copyright: Institute of Software, Chinese Academy of Sciences Beijing ICP No. 05046678-4
Address:4# South Fourth Street, Zhong Guan Cun, Beijing 100190,Postal Code:100190
Phone:010-62562563 Fax:010-62562533 Email:jos@iscas.ac.cn
Technical Support:Beijing Qinyun Technology Development Co., Ltd.

Beijing Public Network Security No. 11040202500063