在 NVIDIA（CUDA、CUBLAS）和 Intel MKL 上直接快速实现 BERT 推理

高度定制化且优化的 BERT 推理，直接运行在 NVIDIA（CUDA、CUBLAS）或 Intel MKL 上，无需 TensorFlow 及其框架开销。

仅支持 BERT（Transformer）模型。

基准测试

环境

• Tesla P4
• 28 核 Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz
• Debian GNU/Linux 8 (jessie)
• gcc (Debian 4.9.2-10+deb8u1) 4.9.2
• CUDA：版本 9.0，V9.0.176
• MKL：2019.0.1.20181227
• tensorflow：1.12.0
• BERT：序列长度 = 32

GPU（cuBERT）

批大小	128（毫秒）	32（毫秒）
tensorflow	255.2	70.0
cuBERT	184.6	54.5

CPU（mklBERT）

批大小	128（毫秒）	1（毫秒）
tensorflow	1504.0	69.9
mklBERT	984.9	24.0

注意：MKL 应在 OMP_NUM_THREADS=? 下运行以控制线程数。其他环境变量及其可能取值包括：

• KMP_BLOCKTIME=0
• KMP_AFFINITY=granularity=fine,verbose,compact,1,0

混合精度

cuBERT 可以利用 NVIDIA Volta 和 Turing 架构 GPU 上的 Tensor Core 和 混合精度计算 来加速推理。我们支持将变量存储为 fp16 格式，并以 fp32 进行计算的混合精度模式。与单精度推理相比，典型精度误差小于 1%，而速度则可提升至两倍以上。

API

API .h 头文件

Pooler

我们支持以下两种池化方法：

• 标准 BERT pooler，定义如下：

with tf.variable_scope("pooler"):
  # 我们通过简单地取第一个 token 对应的隐藏状态来“池化”模型。我们假设该模型已预训练过。
  first_token_tensor = tf.squeeze(self.sequence_output[:, 0:1, :], axis=1)
  self.pooled_output = tf.layers.dense(
    first_token_tensor,
    config.hidden_size,
    activation=tf.tanh,
    kernel_initializer=create_initializer(config.initializer_range))

• 简单平均池化器：

self.pooled_output = tf.reduce_mean(self.sequence_output, axis=1)

输出

支持以下输出类型：

cuBERT_OutputType	python代码
cuBERT_LOGITS	model.get_pooled_output() * output_weights + output_bias
cuBERT_PROBS	probs = tf.nn.softmax(logits, axis=-1)
cuBERT_POOLED_OUTPUT	model.get_pooled_output()
cuBERT_SEQUENCE_OUTPUT	model.get_sequence_output()
cuBERT_EMBEDDING_OUTPUT	model.get_embedding_output()

从源码构建

mkdir build && cd build
# 如果使用 CUDA 构建
cmake -DCMAKE_BUILD_TYPE=Release -DcuBERT_ENABLE_GPU=ON -DCUDA_ARCH_NAME=Common ..
# 或者使用 MKL 构建
cmake -DCMAKE_BUILD_TYPE=Release -DcuBERT_ENABLE_MKL_SUPPORT=ON ..
make -j4

# 安装到 /usr/local
# 如果指定了 -DcuBERT_ENABLE_MKL_SUPPORT=ON，也会同时安装 MKL
sudo make install

如果您想运行 tfBERT_benchmark 进行性能对比，请先从 https://www.tensorflow.org/install/lang_c 安装 TensorFlow C API。

运行单元测试

从 Dropbox 下载 BERT 测试模型 bert_frozen_seq32.pb 和 vocab.txt，并将它们放置在 build 目录下，然后运行 make test 或 ./cuBERT_test。

Python

我们提供了基于 Cython 的简单 Python 封装，可以在 C++ 构建完成后按以下步骤进行构建和安装：

cd python
python setup.py bdist_wheel

# 安装
pip install dist/cuBERT-xxx.whl

# 测试
python cuBERT_test.py

更多详细信息和 Python API 的使用示例，请参阅 cuBERT_test.py。

Java

Java 封装通过 JNA 实现。在安装 Maven 并完成 C++ 构建后，可以按以下步骤进行构建：

cd java
mvn clean package # -DskipTests

使用 Java JAR 文件时，如果 libcuBERT.so 未安装到系统路径中，需要指定 jna.library.path 为该库的路径。此外，应在 JVM 启动脚本中设置 jna.encoding=UTF8。

更多关于 Java API 的使用和示例，请参阅 ModelTest.java。

安装

预编译的 Python 二进制包（目前仅适用于 Linux 上的 MKL）可以按以下步骤安装：

• 下载并安装 MKL 到系统路径。

• 下载 wheel 包并运行 pip install cuBERT-xxx-linux_x86_64.whl。

• 运行 python -c 'import libcubert' 以验证安装是否成功。

依赖项

Protobuf

cuBERT 使用 protobuf-c 构建，以避免与 TensorFlow 的 protobuf 版本和代码发生冲突。

CUDA

不同版本的 CUDA 编译的库之间不兼容。

MKL

MKL 是动态链接的。我们在执行 sudo make install 时会同时安装 cuBERT 和 MKL。

多线程

我们假设 cuBERT 的典型用法是用于在线服务，在这种情况下，需要尽可能快速地处理不同批大小的并发请求。因此，在纯 CPU 环境下，吞吐量和延迟之间需要取得平衡。

由于原生的 Bert 类 因其内部计算缓冲区而不具备线程安全性，因此我们编写了一个包装类 BertM 类，用于为不同的 Bert 实例加锁以确保线程安全。BertM 会以轮询方式选择一个底层的 Bert 实例，因此对同一 Bert 实例的后续请求可能会被其对应的锁阻塞排队。

GPU

每个 Bert 实例都部署在一个 GPU 卡上。最大并发请求数等于一台机器上可用的 GPU 数量，可以通过设置 CUDA_VISIBLE_DEVICES 来控制。

CPU

对于纯 CPU 环境，情况比 GPU 更为复杂。并行性可以分为两个层次：

1. 请求级并行：在多线程的在线服务中，并发请求会竞争 CPU 资源。如果服务器是单线程的（例如某些 Python 实现的服务器），则处理起来会简单得多。

2. 操作级并行：矩阵运算由 OpenMP 和 MKL 进行并行化。最大并行度受 OMP_NUM_THREADS、MKL_NUM_THREADS 以及其他多个环境变量的控制。我们建议用户先阅读 在多线程应用程序中使用线程化的 Intel® MKL 和 从多线程应用程序调用 Intel MKL 例程的推荐设置。

因此，我们引入了 CUBERT_NUM_CPU_MODELS，以便更好地控制请求级并行度。该变量指定了在 CPU/内存上创建的 Bert 实例数量，其作用类似于 GPU 上的 CUDA_VISIBLE_DEVICES。

• 如果您的 CPU 核心数量有限（例如旧式 CPU、桌面 CPU 或 Docker 容器中），则无需设置 CUBERT_NUM_CPU_MODELS。例如，在 4 核 CPU 上，将请求级并行度设为 1、操作级并行度设为 4，通常就能取得不错的效果。

• 但如果您拥有大量 CPU 核心（如 40 核），则可以尝试将请求级并行度设为 5、操作级并行度设为 8。

总之，OMP_NUM_THREADS 或 MKL_NUM_THREADS 决定单个模型可使用的线程数，而 CUBERT_NUM_CPU_MODELS 则决定总的模型数量。

再次强调，每次请求的延迟和整体吞吐量需要权衡，这与模型的序列长度、批大小、CPU 核心数、服务器 QPS 以及许多其他因素有关。您应当进行充分的基准测试，以找到最佳的平衡点。祝您成功！

作者

• fanliwen
• wangruixin
• fangkuan
• sunxian

cuBERT 快速上手指南

cuBERT 是一个高度定制和优化的 BERT 推理引擎，直接基于 NVIDIA CUDA/CUBLAS 或 Intel MKL 实现，无需依赖 TensorFlow 及其框架开销，显著提升推理速度并降低延迟。仅支持 BERT (Transformer) 模型。

环境准备

系统要求

• 操作系统: Linux (推荐 Debian/Ubuntu)
• 编译器: GCC 4.9+
• 硬件加速:
  • GPU 方案: NVIDIA GPU (支持 CUDA 9.0+)，Volta/Turing 架构显卡可开启混合精度加速。
  • CPU 方案: Intel CPU (需安装 Intel MKL)。

前置依赖

根据选择的加速后端，需安装以下依赖：

1. 通用依赖:

   • protobuf-c (用于避免与 TensorFlow 的 protobuf 版本冲突)
   • cmake, make

2. GPU 方案 (cuBERT):

   • NVIDIA CUDA Toolkit (版本需与编译环境一致，如 9.0)
   • CUBLAS

3. CPU 方案 (mklBERT):

   • Intel MKL (Math Kernel Library)
   • OpenMP

注意: 若需进行性能对比测试，可额外安装 TensorFlow C API，但运行 cuBERT 本身不需要。

安装步骤

方式一：源码编译（推荐）

此方法可同时安装 cuBERT 及所需的 MKL 库。

mkdir build && cd build

# 选项 A: 编译 GPU 版本 (CUDA)
cmake -DCMAKE_BUILD_TYPE=Release -DcuBERT_ENABLE_GPU=ON -DCUDA_ARCH_NAME=Common ..

# 选项 B: 编译 CPU 版本 (MKL)
# cmake -DCMAKE_BUILD_TYPE=Release -DcuBERT_ENABLE_MKL_SUPPORT=ON ..

make -j4

# 安装到系统目录 (/usr/local)，若选择 MKL 方案也会一并安装 MKL
sudo make install

方式二：Python 预编译包安装 (仅限 Linux + MKL)

如果你只需要在 Python 中使用 CPU 版本：

1. 下载并安装 Intel MKL 到系统路径。
2. 下载对应的 wheel 包 (cuBERT-xxx-linux_x86_64.whl)。
3. 执行安装：

pip install cuBERT-xxx-linux_x86_64.whl

4. 验证安装：

python -c 'import libcubert'

语言绑定构建 (可选)

• Python Wrapper:

  cd python
  python setup.py bdist_wheel
  pip install dist/cuBERT-xxx.whl
  

• Java Wrapper (需 Maven):

  cd java
  mvn clean package
  # 运行时需指定 -Djna.library.path 指向 libcuBERT.so
  

基本使用

1. 准备模型文件

下载冻结的 BERT 模型文件 (bert_frozen_seq32.pb) 和词表 (vocab.txt)。

• 下载地址：Dropbox 链接
• 将文件放置在构建目录或代码可访问的路径下。

2. Python 使用示例

cuBERT 提供了简洁的 Cython 封装。以下是加载模型并进行推理的最小化示例：

import cubert

# 初始化模型
# model_path: .pb 模型文件路径
# vocab_path: vocab.txt 路径
# seq_length: 序列长度 (需与模型训练时一致，如 32)
# output_type: 输出类型 (LOGITS, PROBS, POOLED_OUTPUT 等)
model = cubert.Bert(
    model_path="bert_frozen_seq32.pb",
    vocab_path="vocab.txt",
    seq_length=32,
    output_type=cubert.cuBERT_PROBS 
)

# 准备输入数据 (列表形式)
texts = ["你好，世界", "cuBERT 推理速度很快"]

# 执行推理
# 返回结果为 numpy 数组
results = model.predict(texts)

print(results)

3. 多线程与并发配置 (重要)

cuBERT 专为在线服务设计，针对多线程环境做了特殊优化：

• GPU 环境: 每个 Bert 实例占用一张显卡。最大并发数等于可用显卡数。可通过 CUDA_VISIBLE_DEVICES 控制。

• CPU 环境: 需平衡“请求级并行”与“运算级并行”。

  • OMP_NUM_THREADS / MKL_NUM_THREADS: 控制单个模型内部矩阵运算使用的线程数。
  • CUBERT_NUM_CPU_MODELS: 控制创建的模型实例总数（即请求级并行度）。

  配置建议:

  • 少核环境 (如 4 核): 设置 CUBERT_NUM_CPU_MODELS=1, OMP_NUM_THREADS=4。
  • 多核环境 (如 40 核): 尝试设置 CUBERT_NUM_CPU_MODELS=5, OMP_NUM_THREADS=8。

  请根据实际压测结果调整这两个参数以获得最佳吞吐量和延迟平衡。

4. 混合精度加速 (GPU)

在 NVIDIA Volta 或 Turing 架构显卡上，cuBERT 自动支持 Tensor Core 混合精度计算（变量存为 fp16，计算用 fp32），通常可获得 2 倍以上加速，精度损失小于 1%。无需额外代码配置，编译时确保开启 CUDA 支持即可。