ByteTransformer - 为BERT类Transformer优化的高性能推理库

ByteTransformer 项目介绍

ByteTransformer 是一个高性能的推理库，专用于优化基于 BERT 的 Transformer 模型在 NVIDIA GPU 上的推理效率。它为各种变长和定长的 Transformer 提供支持，并具备多项优异的性能优化特性。

项目特色

多语言 API：ByteTransformer 提供了 Python 和 C++ 的 API，用户可以通过 PyTorch 插件轻松地用几行 Python 代码提高 Transformer 的推理效率。
支持多种输入长度：无论是固定长度还是可变长度的 Transformer 模型，ByteTransformer 都能提供支持。
专业优化算法：它利用架构感知优化技术，从头到尾对 BERT 常规操作进行了无填充算法优化，包括 QKV 解码、softmax、前馈网络、激活、层归一化和多头注意力等。

ByteTransformer 已大规模部署于字节跳动的内部 Transformer 推理服务系统中，显著提升了定长和变长输入的推理性能，并在 IEEE IPDPS 2023 进行了技术详述。

项目引用

如果使用了 ByteTransformer 库，建议引用其研究论文：

@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}
}

性能表现与对比

ByteTransformer 在 A100 GPU 上与 PyTorch、TensorFlow、FasterTransformer 和 DeepSpeed 的性能进行了对比测试，证明其在标准 BERT 批处理大小和序列长度下拥有显著的执行时间优势。例如，在批处理大小为 1 和 16 的对比测试中，ByteTransformer 均表现出了最快的执行时间，明显优于其它库。

支持的模型

当前支持的模型仅限于标准 BERT Transformer 编码器。

环境要求

字节转换器需要满足以下环境条件：

CUDA 版本为 11.6
CMake 版本 >= 3.13
PyTorch 版本 >= 1.8
GPU 计算能力为 7.0 (V100) / 7.5 (T4) 或 7.0 (A100)
Python 版本 >= 3.7

在 A100 结合 CUDA 11.6、PyTorch 1.13.0+cu116 以及 Python 3.9.16 上进行了合适的测试。

从源码构建

可以通过以下命令从源码构建 ByteTransformer：

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 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

生成测试数据后，可以执行：

./bin/bert_transformer_test 16 64 12 64

在 Python 中进行单位测试

使用与 C++ 测试相同的脚本进行 PyTorch 插件测试，只需省去 --export_data 标志：

# 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

基准测试

可以通过以下命令执行基准测试：

cd build
../benchmark/bert_bench.sh

ByteTransformer 是提升 Transformer 推理效率的有效工具，其出色的性能和灵活性使得它在实际应用中具有很大的优势。