Easy and Efficient Transformer: Scalable Inference Solution For Large NLP Model (2022)

Abstract: Recently, large-scale transformer-based models have been proven to be effective over various tasks across many domains. Nevertheless, applying them in industrial production requires tedious and heavy works to reduce inference costs. To fill such a gap, we introduce a scalable inference solution: Easy and Efficient Transformer (EET) , including a series of transformer inference optimization at the algorithm and implementation levels. First, we design highly optimized kernels for long inputs and large hidden sizes. Second, we propose a flexible CUDA memory manager to reduce the memory footprint when deploying a large model. Compared with the state-of-the-art transformer inference library (Faster Transformer v4.0), EET can achieve an average of 1.40-4.20x speedup on the transformer decoder layer with an A100 GPU.

## introduction in recent years, transformer-based models have achieved impressive performance across variant domains, such as natural language processing @xcite @xcite , computer vision @xcite and speech processing @xcite . the scaling law proposed by @xcite indicates that the validation ppl of a neural language model scales as a power-law with model sizes, dataset sizes, and the amount of training computation. such law is corroborated empirically by many following works @xcite . however, the mega-sized models are notoriously expensive for deployment in the industry. for example, gpt-2 medium model (700m parameters @xcite ) spends up to 10s to generate 512 tokens given a prompt with the length of * equal contribution † corresponding author 512 on an rtx 2080ti gpu, which is not allowed in the industrial application. multiple approaches have been proposed to solve such problems, including knowledge distillation @xcite , model pruning @xcite , and quantization @xcite . apart from these works, much attention has also been paid to optimizing cuda implementation of a transformer layer for better hardware utilization. previous works (e.g.: tensorrt (nvidia, 2021b), light-seq @xcite and faster transformer (ft) (nvidia, 2021a)) have implemented many optimization techniques, including kernels fusion, gemm optimization, quantization, etc. however, these works still have several limitations. tensorrt only contains the multi-head attention(mha) operation, lacking a complete transformer model. light-seq cannot support the model hidden size and input sequence length over 1024. ft contains some performance flaws which need to be improved. in this paper, we propose a novel transformer inference acceleration library, easy and efficient transformer (eet) . first, we implement custom cuda kernels to avoid explicit matrix addition of attention and padding masks with attention weights. as a result, the attention mask matrix is no longer required, while ft spends overhead to initialize an attention mask on the cpu and push it to cuda. in addition, compared with ft, padding masks are no longer needed in computation, leading to additional performance improvement. second, we propose a new method, thread block folding, to extend all kernels to support a larger model size up to 12288 and a longer sequence up to 4096. for ft, it directly assigns the thread number in a cuda block, which may hurt the parallel efficiency. finally, we design a dynamic cuda memory management mechanism to reduce the cuda memory occupation for the same model size, while ft needs to manually allocate memory usage. to compare eet with fairseq, foot_0 lightseq and ft. in our experiments, eet achieves about 4.48-20.27x and 4.30-27.43x speedup over fairseq on 2080ti and a100 respectively. when we set the model size to 768 and 1024 on 2080ti, eet makes 0.82-2.46x speedup over lightseq. compared to ft(v3.1), eet achieves about 1.21-6.30x and 1.62-8.16x speedup on 2080ti and a100 respectively. compared to ft(v4.0), eet achieves about 1.40-4.20x speedup on a100. the remarkable experimental results corroborate the effectiveness of our eet. to remove previously mentioned masks in computation, we redesign the kernels and call the mechanism mask fusion. to extend all the kernels to support the model size or sequence length greater than 1024, we improve the cuda thread structure and call the method as thread block folding. next, we describe these two methods in detail. ## dynamic memory manager the inference is much more sensitive to latency compared to training. as a result, model parallelism @xcite and pipeline parallelism @xcite are undesirable for inference. their communication overhead introduced by tensor slicing or layer split is significant even with the support of nvlink and gpudirect. to reduce the latency and hardware requirements for online service, minimizing the memory footprint is of exceptional value when loading very large models. thus we propose a dynamic memory management strategy for this issue. except for the model weights, the memory footprint includes the caches and the buffers. it is hard to reduce the memory footprint of weights because they are inherent to the model. similarly, the k/v caches for mha are also hard to compress because they are pre-allocated to avoid runtime memory requests, the size of which depends on the model size, maximum batch size, and maximum sequence length. whereas the activation cache and the buffers used to store the operator's results are compressible. hence our dynamic memory management strategy mainly focuses on the activation caches and the operation result buffers. ## experiments unlike fairseq and ft(v4.0), lightseq only supports model sizes that are smaller than 1024, we also make a comparison here as a supplement. figure 7 shows that we make 0.82-2.46x speedup when we set the model size to 768 and 1024. ## conclusion this paper comprehensively describes a series of optimization techniques for transformer inference acceleration exploiting both algorithmic and gpu hardware features. these techniques are packed into the eet, a library dedicated to inference acceleration for large transformer-based models and long input lengths. eet has a 1.40-4.42x speedup for the gpt-2 layer and a 0.99-1.27x speedup for the bert layer compared to the state-of-the-art transformer inference library ft. to make eet easier to apply to a specific task, we provide operation level and model level api, meanwhile integrating web service with dynamic batching. we will continue to improve and keep it up-to-date.