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

TL;DR

ByteTransformer introduces a padding-free, architecture-aware optimization for transformers that significantly accelerates variable-length input processing, outperforming existing frameworks on NVIDIA GPUs.

Contribution

The paper presents a novel padding-free algorithm and architecture optimizations for transformers, especially Multi-Head Attention, to improve performance on variable-length sequences.

Findings

Fused MHA outperforms PyTorch by 6.13x.

ByteTransformer surpasses state-of-the-art frameworks by up to 138%.

Optimization methods are applicable to various BERT-like models.

Abstract

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.