TurboTransformers：一款在 CPU 和 GPU 上运行的快速且易用的 Transformer 推理运行时

为你的推理引擎装上涡轮增压，让 Transformer 服务更加快速！

微信 AI 开源了 TurboTransformers，它具有以下特点：

1. 同时支持 Transformer 的编码器和解码器。
2. 支持可变长度输入。无需耗时的离线调优，你可以实时调整批大小和序列长度。
3. 出色的 CPU/GPU 性能。
4. 完美的易用性。TurboTransformers 提供 Python 和 C++ API。
5. 智能批处理。针对不同长度的请求批次，最小化零填充开销。 它可以作为 PyTorch 的插件使用，只需添加几行 Python 代码即可实现端到端加速。

TurboTransformers 已经应用于腾讯内部多个在线 BERT 服务场景。例如，它使微信 FAQ 服务加速 1.88 倍，公有云情感分析服务加速 2.11 倍，QQ 推荐系统加速 13.6 倍。此外，它还被用于构建聊天、搜索和推荐等服务。

下表是 TurboTransformers 与相关工作的对比：

相关工作	性能	是否需要预处理	可变长度	使用难度
pytorch JIT (CPU)	快	是	否	难
TensorRT (GPU)	快	是	否	难
tf-Faster Transformers (GPU)	快	是	否	难
ONNX-runtime (CPU/GPU)	快/快	否	是	中等
tensorflow-1.x (CPU/GPU)	慢/中等	是	否	易
pytorch (CPU/GPU)	中等/中等	否	是	易
turbo-transformers (CPU/GPU)	最快/最快	否	是	易

支持的模型

我们目前支持以下 Transformer 模型。

• BERT [Python] [C++]
• ALBERT [Python]
• Roberta [Python]
• Transformer 解码器 [Python]
• GPT2 [Python]

用两行 Python 代码加速 BERT 推理

import torch
import transformers
import turbo_transformers

if __name__ == "__main__":
    turbo_transformers.set_num_threads(4)
    torch.set_num_threads(4)
    model_id = "bert-base-uncased"
    model = transformers.BertModel.from_pretrained(model_id)
    model.eval()
    cfg = model.config

    input_ids = torch.tensor(
        ([12166, 10699, 16752, 4454], [5342, 16471, 817, 16022]),
        dtype=torch.long)
    position_ids = torch.tensor(([1, 0, 0, 0], [1, 1, 1, 0]), dtype=torch.long)
    segment_ids = torch.tensor(([1, 1, 1, 0], [1, 0, 0, 0]), dtype=torch.long)
    torch.set_grad_enabled(False)
    torch_res = model(
        input_ids, position_ids=position_ids, token_type_ids=segment_ids
    )  # sequence_output, pooled_output, (hidden_states), (attentions)
    torch_seqence_output = torch_res[0][:, 0, :]
    tt_model = turbo_transformers.BertModel.from_torch(model)
    res = tt_model(
        input_ids, position_ids=position_ids,
        token_type_ids=segment_ids)  # pooled_output, sequence_output
    tt_seqence_output = res[0]

安装

请注意，构建脚本仅适用于特定的操作系统和软件（PyTorch、OpenNMT、transformers 等）版本。请根据你的需求进行调整。

CPU

git clone https://github.com/Tencent/TurboTransformers --recursive

1. 在你的机器上构建 Docker 镜像和容器。

sh tools/build_docker_cpu.sh
# 可选：如果你想在基准测试中比较 onnxrt-mkldnn 的性能，需要将 BUILD_TYPE 设置为 dev，以便将 onnxruntime 编译到 Docker 镜像中，如下所示
env BUILD_TYPE=dev sh tools/build_docker_cpu.sh
docker run -it --rm --name=turbort -v $PWD:/workspace your_image_name /bin/bash

2. 在 Docker 中安装 Turbo

方法 1：我想进行单元测试

cd /workspace
sh tools/build_and_run_unittests.sh $PWD -DWITH_GPU=OFF
# 你可以通过修改 CMakeList.txt 中的这一行在 Openblas 和 MKL 之间切换
# set(BLAS_PROVIDER "mkl" CACHE STRING "设置 BLAS 提供库，在 [openblas, mkl, blis] 中选择")

方法 2：我不想进行单元测试

cd /workspace
mkdir -p build && cd build
cmake .. -DWITH_GPU=OFF
make -j 4
pip install `find . -name *whl`

3. 在 Docker 中运行基准测试（可选），并与 PyTorch、torch-JIT、onnxruntime 进行比较

cd benchmark
bash run_benchmark.sh

4. 在 Docker 中安装 Conda 包（可选）

sh tool/build_conda_package.sh
# Conda 包将位于 /workspace/dist/*.tar.bz2
# 当在该容器之外的其他环境中使用 turbo_transformers 时：conda install your_root_path/dist/*.tar.bz2

我们还准备了一个包含 CPU 版 TurboTransformers 以及其他相关工具（如 onnxrt v1.2.0 和 pytorch-jit）的 Docker 镜像，托管在 Docker Hub 上

docker pull thufeifeibear/turbo_transformers_cpu:latest

GPU

git clone https://github.com/Tencent/TurboTransformers --recursive

1. 在你的机器上构建 Docker 镜像和容器。

# 你可以修改脚本中的环境变量来指定 CUDA 版本和操作系统版本
sh tools/build_docker_gpu.sh $PWD
nvidia-docker run --gpus all --net=host --rm -it -v $PWD:/workspace -v /etc/passwd:/etc/passwd --name=your_container_name REPOSITORY:TAG
# 例如：nvidia-docker run --gpus all --net=host --rm -it -v $PWD:/workspace -v /etc/passwd:/etc/passwd --name=turbo_gpu_env thufeifeibear:0.1.1-cuda9.0-ubuntu16.04-gpu-dev

2. 在 Docker 中安装 pip 包并进行单元测试

cd /workspace
sh tools/build_and_run_unittests.sh $PWD -DWITH_GPU=ON

3. 在 Docker 容器中运行基准测试（可选），并与 PyTorch 进行比较

cd benchmark
bash gpu_run_benchmark.sh

我们还准备了一个包含 GPU 版 TurboTransformers 的 Docker 镜像。

docker pull thufeifeibear/turbo_transformers_gpu:latest

使用 Tensor Core（FP16）

Tensor Core 可以加速 GPU 上的计算。默认情况下，TurboTransformers 中它是关闭的。如果你想启用它，在编译代码之前，请在 CMakeLists.txt 中设置选项 WITH_MODULE_BENCHMAKR ON：

option(WITH_TENSOR_CORE     "使用 Tensor core 加速"     ON)

使用方法

TurboTransformers 提供 C++ 和 Python API 接口。我们致力于适应各种线上环境，以降低用户的开发难度。

预训练模型加载

使用 Turbo 的第一步是加载预训练模型。我们提供了加载 huggingface/transformers 中 PyTorch 和 TensorFlow 预训练模型的方法。具体转换方式是使用 ./tools 目录下的相应脚本将预训练模型转换为 npz 格式文件，然后通过 Turbo 的 C++ 或 Python 接口加载该 npz 文件。特别地，考虑到大多数预训练模型采用 PyTorch 格式并通常与 Python 一起使用，我们还提供了一种直接在 Python 中调用 PyTorch 模型的快捷方式。

API 接口

Python API

请参考 ./example/python 中支持模型的示例。 TurboNLP/Translate-Demo 展示了在翻译任务中应用 TurboTransformer 的演示。 由于 BERT 加速的用户通常需要针对特定任务进行自定义后处理，我们提供了一个编写序列分类应用的示例。

C++ API

请参阅 ./example/cpp 中的示例。 我们的示例提供了 GPU 以及两种 CPU 多线程调用方式：一种是使用多线程执行单次 BERT 推理；另一种则是多个 BERT 推理任务，每个任务由一个线程独立完成。 用户可以通过 add_subdirectory 将 turbo-transformers 链接到自己的代码中。

智能批处理（最小化批处理中的零填充开销）

通常情况下，当我们将一批长度不同的请求输入到 BERT 模型进行推理时，需要对这些请求进行零填充，以使它们具有相同的长度。例如，如果服务的请求列表长度分别为 (100, 10, 50)，则需要在预处理阶段将其填充为 (100, 100, 100)。这样一来，后两个序列分别有 90% 和 50% 的计算量被浪费掉了。正如 Effective Transformer 所指出的那样，实际上并不需要对输入张量进行零填充。相反，只需对多头注意力机制内部的批量矩阵乘法操作进行填充即可，而这部分填充仅占整个 BERT 计算的一小部分。因此，大部分矩阵乘法运算都可以在不进行零填充的情况下完成。Turbo 提供了一个名为 BertModelSmartBatch 的模型，其中包含了智能批处理技术。相关示例可在 ./example/python/bert_smart_pad.py 中找到。

如何贡献新模型

如何找出代码中的热点？

如何添加新层？

待办事项

目前（2020年6月），在不久的将来，我们将增加对低精度模型的支持（CPU int8、GPU FP16）。 期待您的贡献！

许可证

BSD 3-Clause 许可证

已知问题

1. Turbo Transformers 的结果可能与 PyTorch 的结果在小数点后两位之后有所不同。这种差异主要来自 BERT 的输出层。我们使用了一种近似的 GELU 算法，这可能与 PyTorch 的实现有所区别。
2. Turbo 和 PyTorch 共享同一个 MKL 库。PyTorch 1.5.0 版本的 MKL 可能在 Turbo 中运行较慢，具体原因尚待确认。将 PyTorch 版本降级到 1.1.0 可以提升 Turbo 的性能。
3. onnxruntime-cpu==1.4.0 和 onnxruntime-gpu==1.3.0 无法同时工作。

历史

1. 2021年1月 v0.6.0，TurboTransformers 支持智能批处理。
2. 2020年7月 v0.4.0，TurboTransformers 使用 onnxruntime 作为 CPU 后端，支持 GPT2，并新增了量化 BERT。
3. 2020年7月 v0.3.1，TurboTransformers 在 CPU/GPU 上增加了对 ALbert 和 Roberta 的支持。
4. 2020年6月 v0.3.0，TurboTransformers 在 CPU/GPU 上增加了对 Transformer 解码器的支持。
5. 2020年6月 v0.2.1，TurboTransformers 添加了 BLIS 作为 BLAS 提供者选项，在 AMD CPU 上表现更佳。
6. 2020年4月 v0.0.1，TurboTransformers 正式发布，并实现了 CPU/GPU 上最先进的 BERT 推理速度。

引用我们

如果您在研究论文中使用了 TurboTransformers，请引用以下文献：

@inproceedings{fang2021turbotransformers,
  title={TurboTransformers: an efficient GPU serving system for transformer models},
  author={Fang, Jiarui and Yu, Yang and Zhao, Chengduo and Zhou, Jie},
  booktitle={Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming},
  pages={389--402},
  year={2021}
}

该论文的相关成果可在 ppopp21_artifact_centos 分支中找到。

联系我们

虽然我们建议您通过 GitHub Issues 提出问题，但您也可以加入我们的 Turbo 用户群。

1. 扫描此 二维码，将我们的联系人添加为您的微信好友。
2. QQ 群名称：TurboTransformers，群号：1109315167。

TurboTransformers 快速上手指南

TurboTransformers 是腾讯微信 AI 开源的高性能 Transformer 推理运行时，支持 CPU 和 GPU。它无需耗时的离线调优，支持变长输入和智能批处理（Smart Batching），可显著提升 BERT、GPT2 等模型的推理速度（最高可达 13.6 倍加速），并提供易用的 Python 和 C++ 接口。

环境准备

• 操作系统: Linux (推荐 Ubuntu 16.04/18.04/20.04)
• 编译器: GCC 5.4+ 或更高版本
• 依赖软件:
  • CMake 3.10+
  • Python 3.6+
  • PyTorch (建议 1.1.0 - 1.7.0，部分版本 MKL 兼容性更佳)
  • CUDA & cuDNN (仅 GPU 版本需要，推荐 CUDA 9.0+/10.0+)
• 硬件要求:
  • CPU: 支持 AVX2 指令集
  • GPU: NVIDIA GPU (支持 Tensor Core 的显卡可获得 FP16 加速)

提示: 为了简化环境配置，官方强烈推荐使用提供的 Docker 镜像进行开发和部署。

安装步骤

方式一：使用 Docker（推荐）

这是最快捷的方式，已预装所有依赖。

1. CPU 版本

# 拉取镜像
docker pull thufeifeibear/turbo_transformers_cpu:latest

# 运行容器
docker run -it --rm --name=turbort -v $PWD:/workspace thufeifeibear/turbo_transformers_cpu:latest /bin/bash

2. GPU 版本

# 拉取镜像
docker pull thufeifeibear/turbo_transformers_gpu:latest

# 运行容器 (需确保宿主机安装了 nvidia-docker)
nvidia-docker run --gpus all --net=host --rm -it -v $PWD:/workspace --name=turbo_gpu_env thufeifeibear/turbo_transformers_gpu:latest /bin/bash

方式二：源码编译安装

如果你需要在本地环境中直接安装，请按以下步骤操作。

1. 克隆代码

git clone https://github.com/Tencent/TurboTransformers --recursive
cd TurboTransformers

2. 编译与安装 (CPU)

mkdir -p build && cd build
cmake .. -DWITH_GPU=OFF
make -j 4
# 安装生成的 wheel 包
pip install `find . -name *whl`

3. 编译与安装 (GPU)

mkdir -p build && cd build
cmake .. -DWITH_GPU=ON
make -j 4
pip install `find . -name *whl`

注意: 若需启用 GPU Tensor Core (FP16) 加速，请在 cmake 命令中添加 -DWITH_TENSOR_CORE=ON。

基本使用

TurboTransformers 可以无缝集成到现有的 PyTorch 代码中。只需几行代码即可将 Hugging Face transformers 加载的模型转换为高性能推理模型。

以下是一个加速 BERT 推理的最小化示例：

import torch
import transformers
import turbo_transformers

if __name__ == "__main__":
    # 设置线程数以获得最佳性能
    turbo_transformers.set_num_threads(4)
    torch.set_num_threads(4)
    
    # 1. 加载原生 PyTorch 模型
    model_id = "bert-base-uncased"
    model = transformers.BertModel.from_pretrained(model_id)
    model.eval()
    
    # 准备输入数据
    input_ids = torch.tensor(
        ([12166, 10699, 16752, 4454], [5342, 16471, 817, 16022]),
        dtype=torch.long)
    position_ids = torch.tensor(([1, 0, 0, 0], [1, 1, 1, 0]), dtype=torch.long)
    segment_ids = torch.tensor(([1, 1, 1, 0], [1, 0, 0, 0]), dtype=torch.long)
    
    torch.set_grad_enabled(False)
    
    # (可选) 原生 PyTorch 推理
    # torch_res = model(input_ids, position_ids=position_ids, token_type_ids=segment_ids)
    
    # 2. 转换为 TurboTransformers 模型 (核心步骤)
    tt_model = turbo_transformers.BertModel.from_torch(model)
    
    # 3. 执行高速推理
    res = tt_model(
        input_ids, position_ids=position_ids,
        token_type_ids=segment_ids)
    
    # 获取结果 (pooled_output, sequence_output)
    tt_seqence_output = res[0]
    print("Inference completed with TurboTransformers.")

进阶特性：智能批处理 (Smart Batching)

针对不同长度的输入序列，TurboTransformers 提供了 BertModelSmartBatch 类，可自动最小化零填充（Zero-Padding）带来的计算浪费，进一步提升吞吐量。具体用法可参考官方示例 ./example/python/bert_smart_pad.py。