ByteTransformer:在 NVIDIA GPU 上优化的 BERT Transformer 推理
简介
ByteTransformer 是一个高性能的 BERT 类 transformer 推理库,具有以下特点:
- 提供 Python 和 C++ API,PyTorch 插件允许用户只需几行 Python 代码即可增强 transformer 推理。
- 支持定长和变长 transformer。
- 包含对 BERT 例程中无填充算法的端到端架构感知优化,包括 QKV 编码、softmax、前馈网络、激活、层归一化和多头注意力。
ByteTransformer 已广泛部署于 ByteDance 内部的 transformer 推理服务系统中,与其他 transformer 实现相比,提供了更出色的性能,支持定长和变长输入。技术细节已在 IEEE IPDPS 2023 上发表。
引用我们
如果您使用了我们的库,请引用我们的研究论文。
@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}
}
性能和加速
我们在 A100 GPU 上将 ByteTransformer 与 PyTorch、TensorFlow、FasterTransformer 和 DeepSpeed 进行了比较。基准测试脚本可在benchmark/bert_bench.sh中找到。
1. 标准 BERT batch size = 1,平均序列长度 = 0.6 * 最大,执行时间(毫秒):
PyTorch | Tensorflow | FasterTransformer | FasterTransformer 取消填充 | 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. 标准 BERT batch size = 16,平均序列长度 = 0.6 * 最大,执行时间(毫秒):
PyTorch | Tensorflow | FasterTransformer | FasterTransformer 取消填充 | 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 |
支持的模型
目前,该仓库只提供标准 BERT transformer 编码器。
环境要求
- CUDA:11.6
- CMake:>= 3.13
- PyTorch:>= 1.8
- GPU 计算能力:7.0(V100)/ 7.5(T4)或 8.0(A100)
- Python:>= 3.7
测试环境:A100 + CUDA 11.6 + PyTorch 1.13.0+cu116 + Python 3.9.16
从源码构建
从源码构建,请运行以下命令:
git submodule update --init
mkdir build && cd build
cmake -DTORCH_CUDA_ARCH_LIST="8.0" -DDataType=FP16 -DBUILD_THS=ON -DCUDAARCHS="80" ..
make
从单元测试开始
C++ 中的单元测试
生成测试数据,运行以下代码:
cd build
# batch size = 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
这里,16
、64
、12
和 64
分别代表 batch size、序列长度、头数和头大小。--avg_seqlen 32
用于设置平均序列长度,--dtype fp16
设置数据类型,--export_data
导出测试数据。
生成测试数据后(*.in
和 *.out
文件保存在当前目录下),运行以下命令:
./bin/bert_transformer_test 16 64 12 64
这里,参数表示生成测试数据时使用的相同参数。
使用 Python 中的 PyTorch 插件进行单元测试
在 PyTorch 中使用 Python 插件进行单元测试,请使用与 C++ 相同的脚本,但不带 --export_data
标志。在终端中运行以下命令:
# batch size = 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
同样,参数表示生成测试数据时使用的相同参数。
基准测试
cd build
../benchmark/bert_bench.sh