Minimizing transformer inference overhead using controlling element on Shenwei AI accelerator

Abstract

Transformer models have become a cornerstone of various natural language processing (NLP) tasks. However, the substantial computational overhead during the inference remains a significant challenge, limiting their deployment in practical applications. In this study, we address this challenge by minimizing the inference overhead in transformer models using the controlling element on artificial intelligence (AI) accelerators. Our work is anchored by four key contributions. First, we conduct a comprehensive analysis of the overhead composition within the transformer inference process, identifying the primary bottlenecks. Second, we leverage the management processing element (MPE) of the Shenwei AI (SWAI) accelerator, implementing a three-tier scheduling framework that significantly reduces the number of host-device launches to approximately 1/10 000 of the original PyTorch-GPU setup. Third, we introduce a zero-copy memory management technique using segment-page fusion, which significantly reduces memory access latency and improves overall inference efficiency. Finally, we develop a fast model loading method that eliminates redundant computations during model verification and initialization, reducing the total loading time for large models from 22 128.31 ms to 1041.72 ms. Our contributions significantly enhance the optimization of transformer models, enabling more efficient and expedited inference processes on AI accelerators.

1. Introduction

Over the past decade, pre-trained language models based on the transformer architecture (Vaswani et al., 2017) have become a leading paradigm in the domain of natural language processing (NLP). Notable examples of such models include BERT (Devlin et al., 2019), GPT-2 (Radford et al., 2019), LLaMA (Touvron et al., 2023), wav2vec2.0 (Baevski et al., 2020), and whisper (Radford et al., 2023). These transformer models have significantly advanced state-of-the-art models in terms of accuracy, surpassing traditional models. Nonetheless, the computational intensity during the inference phase poses a significant barrier to their integration into real-world applications, which require strict criteria such as low latency, rapid inference capabilities, and cost-effective operational overhead.

There are currently two strategies for enhancing the inference efficiency of transformer models. The first pertains to the optimization of specific operators, such as Softermax (Stevens et al., 2021) and FLASHATTENTION (Dao et al., 2022), and involves reducing accuracy and minimizing input/output (I/O) costs to enhance the operators’ computational speed. The second focuses on the optimization of the model’s runtime during inference, using various techniques such as structured pruning (Kim YJ et al., 2020), knowledge distillation (Fang et al., 2021), operator fusion, memory management (Chen SY et al., 2021), and parallel scheduling strategies (Du et al., 2022). These approaches collectively aim to improve inference efficacy by reducing model size and enhancing overall system performance. However, the intrinsic computational overhead (Ma et al., 2021) associated with transformer inference has received relatively little scholarly attention, despite its critical role and the need for comprehensive research in this area.

To minimize the overhead during the inference process, we propose an innovative approach that uses the hardware control mechanisms of artificial intelligence (AI) accelerators. We begin by conducting a detailed analysis of the transformer model’s inference control process, exploring the components of the runtime overhead through carefully designed experiments. Based on this analysis, we introduce a fast loading method that significantly reduces the overhead caused by loading models with PyTorch-GPU. We also develop a three-tier scheduling framework for interacting with the host and the controlling element on the accelerator, leveraging the accelerator’s scheduling capabilities. To minimize device’s launch expenditures, we embrace a holistic full-graph optimization strategy. Additionally, we implement a zero-copy memory management protocol based on segment-page fusion, which eliminates data transmission-related costs. These optimizations target the reduction of overhead throughout the inference process.