LightSeq: A High Performance Inference Library for Transformers (2010.13887v4)

Abstract: Transformer, BERT and their variants have achieved great success in natural language processing. Since Transformer models are huge in size, serving these models is a challenge for real industrial applications. In this paper, we propose LightSeq, a highly efficient inference library for models in the Transformer family. LightSeq includes a series of GPU optimization techniques to to streamline the computation of neural layers and to reduce memory footprint. LightSeq can easily import models trained using PyTorch and Tensorflow. Experimental results on machine translation benchmarks show that LightSeq achieves up to 14x speedup compared with TensorFlow and 1.4x compared with FasterTransformer, a concurrent CUDA implementation. The code is available at https://github.com/bytedance/lightseq.

LightSeq: A High Performance Inference Library for Transformers

Overview

The paper presents "LightSeq," an inference library specifically designed to optimize the performance of Transformer models, including BERT, GPT, and VAEs, on GPUs. The primary objective of LightSeq is to address the challenges of deploying large Transformer models in real-world applications by enhancing inference efficiency and reducing memory usage.

Key Contributions

The authors introduce several GPU optimization techniques that significantly enhance inference speed. Notable achievements include a 14x speedup compared to TensorFlow and a 1.4x improvement over the FasterTransformer library. LightSeq seamlessly integrates with models trained in PyTorch and TensorFlow, offering a practical solution for industrial applications.

Innovative Features

1. Operation Fusion: LightSeq employs coarse-grained fused GPU kernel functions in place of the fine-grained approach used by frameworks like TensorFlow. This change decreases unnecessary GPU memory I/O and reduces kernel launches by fourfold.
2. Hierarchical Auto Regressive Search (HARS): This method optimizes auto-regressive search tasks by erasing redundant calculations and leveraging a hierarchical strategy to reduce computational inputs size. The approach enhances inference efficiency, particularly in scenarios like beam search.
3. Dynamic GPU Memory Reuse: LightSeq handles variable-length input sequences by pre-defining maximum memory allocations, allowing GPU memory sharing across non-dependent processes. This strategy reduces memory allocation demands without compromising inference speed, achieving an eightfold reduction in memory allocation.

Performance Evaluation

The authors conduct experiments on NVIDIA Tesla P4 and T4 GPUs, demonstrating LightSeq's superiority over existing libraries. In machine translation benchmarks, it outperforms TensorFlow and FasterTransformer, particularly with larger batch sizes and sequences. Text generation tasks, including top-$k$ and top-$p$ sampling, further illustrate LightSeq's enhanced performance, especially in scenarios with smaller batch sizes and shorter sequences.

Implications and Future Directions

LightSeq addresses critical deployment challenges of large Transformer models, significantly narrowing the performance gap necessary for online services. Its ability to integrate efficiently with industry-standard training frameworks and support various model architectures, including multilingual variants, expands its applicability in diverse NLP tasks.

Future research could explore further optimization techniques, potentially incorporating integer arithmetic-only inference and sparse GEMM computations. Such advancements may yield additional speed improvements and efficiency gains in deploying Transformer models at scale.

This work contributes a robust tool for researchers and practitioners seeking to optimize the deployment of complex NLP models in production environments, reflecting an ongoing evolution in optimizing deep learning inference.