大规模语言模型如 BERT，GPT-2 和 XL-Net 在多种自然语言理解任务中表现出了激动人心的准确率。自去年 10 月份 Google 发布 BERT 以来，一直保持为最受欢迎的语言模型之一。然而鲜有文章介绍如何将 BERT 模型进行生产部署，诸如 如何缩短计算延迟并保持准确率以满足用户体验 。本文将填补这一空白。
由于 BERT 使用 12/24 层堆叠多头注意力网络结构，推理需要大量计算，这对需要部署 BERT 模型作为部分实时应用的公司来说是一个巨大挑战。
TensorRT 是 Nvidia 提供的高性能深度学习推理加速库，提供了 CNN, LSTM 等常见模型的加速版本，可以满足高吞吐（Batch Size > 16）或低延迟（Batch Size <= 16）部署需求。 在 Tesla T4 GPU 上可以在 2.2 ms 时间内完成 BERT 推理计算，相比 CPU 计算平台加速超过 17 倍 ！！！
代码已开源：
https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/
01
—
BERT 训练和推理流水线
NLP 研究人员和开发者遇到的最大问题是针对具体 NLP 任务缺少高质量带标签的训练数据。最近的突破是利用大量未标注的文本数据，将 NLP 任务分解为两步：首先利用大量语料构建一个语言模型（预训练/pre-training），然后将语言模型适配到具体任务，增加相对较小的任务相关网络结构，使用监督学习方式训练该部分网络（微调/fine-tuning）。
在这个例子中，使用 BERT 作为高质量语言模型进行预训练，然后微调用作问答系统（QA）。除此之外，还可以微调用作句子分类、情感分析等其他任务。
BERT 预训练过程可以直接用官方训练好的模型权重（如上图左侧所示），也可以从头训练（如上图右侧所示）。当然，你也可以基于官方权重针对自己语料库继续训练。预训练过程计算密集，需要大规模 GPU 计算基础设施加速训练，一般只有大厂才有资源做这个事情。
BERT 微调过程需要使用任务相关训练数据，例如 QA 系统使用训练数据为（段落，问题，回答）三元组。微调过程相比预训练过程对计算量需求大大降低，一般可以在单张 GPU 上短时间内完成。
微调完成后，为了实现推理过程，需要做这几件事：
- 创建 TensorRT 引擎，需要使用微调后的权重和网络定义；
- 利用上面创建的引擎启动 TensorRT Runtime；
- 向 TensorRT Runtime 送入 Passage 和 Question，获取 Answer；
QA 推理完整流程如下图所示：
02
—
BERT 推理优化
BERT 架构基于 Transformer，其中 BERT-base 包含 12 个单元，BERT-large 包含 24 个单元。每个单元包括两个连续残差块，每个残差块后面跟一个 layer normalization（简称 LN）。第一个残差块把其中一个 FC 层换为多头自注意力机制，第二个残差块使用 GELU 激活。单元结构如下图所示：
TensorRT 加速 BERT 主要是 Transformer 单元的优化。基于自定义模块（plugin）将多个计算步骤合并到单一 CUDA kernel 中，节省访存开销。以 GELU 计算为例，TensorFlow 实现代码为：
相应的计算图如下：
我们发现 GELU 激活函数实际是组合了多个 TensorFlow Math OP（Mul, Add, Pow, Tanh 等），这在 GPU 上计算过程中会有多次 global memory 拷贝，势必会降低性能。为了优化这个激活函数，专门编写了 TensorRT plugin，将所有计算步骤包装在一个 CUDA kernel 中。具体代码可以查看 geluPlugin.cu。
https://github.com/NVIDIA/TensorRT/blob/release/5.1/demo/BERT/plugins/geluPlugin.cu
注：这里 GELU 实现均为近似公式计算，参考 GELU 原始论文
gelu ( x )
= a
x *
( 1
tanh ( b
( x
c * x ^ 3 )
))
除了 GELU 之外，LN 和跨层连接也做了合并。具体参考 skipLayerNormPlugin.cu
https://github.com/NVIDIA/TensorRT/blob/release/5.1/demo/BERT/plugins/skipLayerNormPlugin.cu
Transformer 中的自注意力层计算流程如下：
从输入 embedding 通过三个全连接层计算得到 query(Q), key(K) 和 value(V) ，三个全连接层的输入和输出尺寸均为 B x S x (N * H)，其中 B 为 Batch size，S 为 序列长度，N 为注意力 head 数目， H 为隐层单元数目。 将三个全连接层用一个大矩阵乘替换之后，输出尺寸为 B x S x (3 * N * H)。 经过 FC 计算，还有三个矩阵转置计算也可以合并为一个更大的矩阵转置，输出尺寸为 3 x B x N x S x H。 优化之后计算流程变为：
优化之后，Q, K, V 在存储空间中连续排布，有利于更高效地进行后续计算。图中也顺带将 Scaling 和 Softmax 两层做了合并。
代码参考 qkvToContextPlugin.cu
https://github.com/NVIDIA/TensorRT/blob/release/5.1/demo/BERT/plugins/qkvToContextPlugin.cu
03
—
BERT 推理性能评测
BERT 可以应用到在线和离线场景。在线 NLU 应用如对话机器人对延迟非常敏感，每个用户请求需要触发几个模型顺序执行得到返回结果。作为在线服务的总响应时间包括计算时间和输入、输出网络延迟，过长的响应时间会导致用户体验差，从而流失客户。一些实时应用需要语言模型执行时间不超过 10 ms。使用 TensorRT 加速的 BERT 在一块 Tesla T4 GPU 上执行时间仅 2.2 ms，而同样参数下使用 CPU 计算则需要 40 ms，总体性能差距达 17 倍！
CPU 配置：Gold 6240@2.6 GHz, 3.9 GHz Turbo(Cascade Lake), HT Off, 单节点，单 socket，CPU 线程数 18，真实数据，Batch Size =1，Seq Len = 128，nireq = 1，FP32，OpenVINO 2019 R2;
GPU 配置：Gold 6140@2GHz, 3.7 GHz Turbo(Skylake), HT On, 单节点，双 socket，CPU 线程数 72，Tesla T4 16 GB, GPU 驱动版本 418.67(r418_00), BERT-base，Batch Size=1，heads = 12，size per head = 64, 12 layers, seq len = 128, FP16，XLA on，TensorRT 5.1;