ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs

• Published in: Proceedings - IEEE International Parallel and Distributed Processing Symposium pp. 344 - 355
• Main Authors: Zhai, Yujia, Jiang, Chengquan, Wang, Leyuan, Jia, Xiaoying, Zhang, Shang, Chen, Zizhong, Liu, Xin, Zhu, Yibo
• Format: Conference Proceeding
• Language: English
• Published: IEEE 01.05.2023
• Subjects: BERT; Bit error rate; CUTLASS; Deep learning; Distributed processing; Graphics processing units; Large Language Models; Multi-head Attention; Natural Language Processing; NVIDIA GPU; Optimization methods; Technological innovation; Training; Transformer
• ISSN: 1530-2075
• DOI: 10.1109/IPDPS54959.2023.00042

Transformers have become keystone models in natural language processing over the past decade. They have achieved great popularity in deep learning applications, but the increasing sizes of the parameter spaces required by transformer models generate a commensurate need to accelerate performance. Natural language processing problems are also routinely faced with variable-length sequences, as word counts commonly vary among sentences. Existing deep learning frameworks pad variable-length sequences to a maximal length, which adds significant memory and computational overhead. In this paper, we present ByteTransformer, a high-performance transformer boosted for variable-length inputs. We propose a padding-free algorithm that liberates the entire transformer from redundant computations on zero padded tokens. In addition to algorithmic-level optimization, we provide architecture-aware optimizations for transformer functional modules, especially the performance-critical algorithm Multi-Head Attention (MHA). Experimental results on an NVIDIA A100 GPU with variable-length sequence inputs validate that our fused MHA outperforms PyTorch by 6.13x. The end-to-end performance of ByteTransformer for a forward BERT transformer surpasses state-of-the-art transformer frameworks, such as PyTorch JIT, TensorFlow XLA, Tencent TurboTransformer, Microsoft DeepSpeed-Inference and NVIDIA FasterTransformer, by 87%, 131%, 138%, 74% and 55%, respectively. We also demonstrate the general applicability of our optimization methods to other BERT-like models, including ALBERT, DistilBERT, and DeBERTa.