ByteTransformer: Optimized BERT Transformer Inference on NVIDIA GPUs

Introduction

ByteTransformer is a high-performance inference library for BERT-like transformers that offers the following features:

• Provides Python and C++ APIs, with the PyTorch plugin allowing users to enhance transformer inference with just a few lines of Python code.
• Supports both fixed-length and variable-length transformers.
• Includes end-to-end architectural-aware optimizations for the padding-free algorithm on BERT routines, including QKV encoding, softmax, feed forward network, activation, layernorm, and multi-head attention.

ByteTransformer has been widely deployed to improve in-house transformer inference serving systems at ByteDance, delivering superior performance over other transformer implementations for both fixed-length and variable-length inputs. The technical details have been published at IEEE IPDPS 2023.

Cite Us

If you use our library, please cite our research paper.

@article{zhai2022bytetransformer,
  title={ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs},
  author={Zhai, Yujia and Jiang, Chengquan and Wang, Leyuan and Jia, Xiaoying and Zhang, Shang and Chen, Zizhong and Liu, Xin and Zhu, Yibo},
  journal={arXiv preprint arXiv:2210.03052},
  year={2022}
}

Performance and Speedup

We compared ByteTransformer with PyTorch, TensorFlow, FasterTransformer, and DeepSpeed on an A100 GPU. The benchmark script is available in benchmark/bert_bench.sh.

1. Standard BERT batch size = 1, average sequence length = 0.6 * maximal, execution time in millisecond:

| | PyTorch | Tensorflow | FasterTransformer | FasterTransformer with remove padding | DeepSpeed | ByteTransformer | |------|-------------|----------------|-------------------|---------------------------------------|---------------------|-----------------| | 64 | 2.93 | 2.46 | 1.05 | 1.23 | 1.17 | 0.90 | | 128 | 3.18 | 2.6 | 1.10 | 1.43 | 1.28 | 0.97 | | 192 | 3.18 | 2.81 | 1.26 | 1.43 | 1.40 | 1.36 | | 256 | 2.81 | 2.9 | 1.35 | 1.55 | 1.51 | 1.43 | | 320 | 3.11 | 3.24 | 1.63 | 1.66 | 1.84 | 1.69 | | 384 | 2.87 | 3.43 | 1.64 | 1.64 | 1.95 | 1.72 | | 448 | 2.99 | 3.61 | 2.26 | 2.35 | 2.23 | 1.86 | | 512 | 2.89 | 3.74 | 2.28 | 2.43 | 2.37 | 2.00 | | 576 | 2.99 | 4.03 | 2.51 | 2.59 | 2.70 | 2.19 | | 640 | 2.99 | 4.54 | 2.85 | 2.83 | 3.17 | 2.23 | | 704 | 3.21 | 4.67 | 3.16 | 3.44 | 3.32 | 2.47 | | 768 | 3.33 | 4.88 | 3.26 | 3.63 | 3.46 | 2.51 | | 832 | 3.78 | 5.39 | 3.75 | 3.87 | 3.97 | 2.80 | | 896 | 3.86 | 5.81 | 4.08 | 4.95 | 4.37 | 2.86 | | 960 | 4.02 | 6.27 | 4.30 | 5.23 | 4.66 | 3.12 | | 1024 | 4.2 | 6.37 | 4.51 | 4.96 | 4.86 | 3.16 |

2. Standard BERT batch size = 16, average sequence length = 0.6 * maximal, execution time in millisecond:

| | PyTorch | Tensorflow | FasterTransformer | FasterTransformer with remove padding | DeepSpeed | ByteTransformer | |------|-------------|----------------|-------------------|---------------------------------------|---------------------|-----------------| | 64 | 3.2 | 4.57 | 2.24 | 1.93 | 2.81 | 2.09 | | 128 | 4.97 | 6.97 | 3.62 | 3.33 | 4.54 | 3.18 | | 192 | 7.65 | 9.37 | 5.26 | 5.29 | 6.68 | 5.08 | | 256 | 9.56 | 12.17 | 6.77 | 5.49 | 9.03 | 6.85 | | 320 | 13.21 | 15.87 | 8.85 | 6.47 | 12.81 | 7.49 | | 384 | 15.01 | 18.56 | 10.37 | 7.05 | 15.19 | 8.44 | | 448 | 19.06 | 23.01 | 15.97 | 12.54 | 18.83 | 8.89 | | 512 | 21 | 26.03 | 18.03 | 13.79 | 21.55 | 9.22 | | 576 | 24.33 | 31.24 | 21.11 | 17.65 | 26.2 | 10.15 | | 640 | 28.03 | 35.07 | 24.52 | 20.34 | 30.24 | 12.04 | | 704 | 32.33 | 41.43 | 28.94 | 24.52 | 34.65 | 13.55 | | 768 | 35.31 | 44.62 | 32.09 | 28.21 | 37.95 | 16.3 | | 832 | 40.75 | 51.87 | 36.33 | 31.69 | 45.32 | 16.92 | | 896 | 44.47 | 55.65 | 42.17 | 38.05 | 49.48 | 20.67 | | 960 | 49.72 | 63.59 | 47.01 | 42.98 | 55.72 | 23.27 | | 1024 | 53.21 | 65.94 | 50.28 | 45.22 | 59.96 | 24.70 |

Supported Models

Currently, only the standard BERT transformer encoder is available under this repository.

Environment requirements

• CUDA: 11.6
• CMake: >= 3.13
• PyTorch: >= 1.8
• GPU compute capability: 7.0(V100) / 7.5(T4) or 8.0(A100)
• Python: >= 3.7

Tested on: A100 + CUDA 11.6 + PyTorch 1.13.0+cu116 + Python 3.9.16

Building from Source

To build from source, run the following commands:

git submodule update --init
mkdir build && cd build
cmake -DTORCH_CUDA_ARCH_LIST="8.0" -DDataType=FP16 -DBUILD_THS=ON -DCUDAARCHS="80" ..
make

Getting Started with Unit Tests

Unit Tests in C++

To generate test data, run the following code:

cd build
# batch sz = 16, seqlen = 64, head num = 12, head sz = 64, avg seqlen = 32
python3 bert_transformer_test.py 16 64 12 64 --avg_seqlen 32 --dtype fp16 --export_data

Here, 16, 64, 12, and 64 represent batch size, sequence length, number of heads, and head size, respectively. The --avg_seqlen 32 flag is used to set the average sequence length, --dtype fp16 sets the data type, and --export_data exports the test data.

After test data is generated (*.in and *.out files are saved under the current directory), run the following command:

./bin/bert_transformer_test 16 64 12 64

Here, the arguments represent the same parameters as used in generating the test data.

Unit Tests in a PyTorch Plugin in Python

To perform the unit tests in a PyTorch plugin in Python, use the same script as for C++, but without the --export_data flag. Run the following command in the terminal:

# batch sz = 16, seqlen = 64, head num = 12, head sz = 64, avg seqlen = 32
python3 bert_transformer_test.py 16 64 12 64 --avg_seqlen 32 --dtype fp16

Again, the arguments represent the same parameters as used in generating the test data.

Benchmark

cd build
../benchmark/bert_bench.sh