NVIDIA BERT推理解决方案Faster Transformer开源了

Faster Transformer是一个基于CUDA和cuBLAS的Transformer Encoder前向计算实现，其优越的性能将助力于多种BERT的应用场景。

2017 年 12 月 Google 在论文“Attention is All You Need”[1] 中首次提出了 Transformer，将其作为一种通用高效的特征抽取器。至今，Transformer 已经被多种 NLP 模型采用，比如 BERT[2] 以及上月发布重刷其记录的 XLNet[3]，这些模型在多项 NLP 任务中都有突出表现。在 NLP 之外， TTS，ASR 等领域也在逐步采用 Transformer。可以预见，Transformer 这个简洁有效的网络结构会像 CNN 和 RNN 一样被广泛采用。虽然 Transformer 在多种场景下都有优秀的表现，但是在推理部署阶段，其计算性能却受到了巨大的挑战：以 BERT 为原型的多层 Transformer 模型，其性能常常难以满足在线业务对于低延迟（保证服务质量）和高吞吐（考虑成本）的要求。以 BERT-BASE 为例，超过 90% 的计算时间消耗在 12 层 Transformer 的前向计算上。因此，一个高效的 Transformer 前向计算方案，既可以为在线业务带来降本增效的作用，也有利于以 Transformer 结构为核心的各类网络在更多实际工业场景中落地。本文将介绍 NVIDIA GPU 计算专家团队针对 Transformer 推理提出的性能优化方案：Faster Transformer。

Faster Transformer 是一个 BERT Transformer 单层前向计算的高效实现，其代码简洁明了，后续可以通过简单修改支持多种 Transformer 结构。目前优化集中在编码器（encoder）的前向计算（解码器 decoder 开发在后续特性规划中）。底层由 CUDA 和 cuBLAS 实现，支持 FP16 和 FP32 两种计算模式，其中 FP16 可以充分利用 Volta 和 Turing 架构 GPU 上的 Tensor Core 计算单元。

Faster Transformer 共接收 4 个输入参数。首先是 attention head 的数量以及每个 head 的维度。这两个参数是决定 Transformer 网络结构的关键参数。这两个参数的动态传入，可以保证 Faster Transformer 既支持标准的 BERT-BASE（12 head x 64维），也支持裁剪过的模型（例如，4 head x 32 维），或者其他各式专门定制化的模型。其余两个参数是 Batch Size 和句子最大长度。出于性能考虑，目前句子最大长度固定为最常用的 32，64 和 128 三种，未来会支持任意长度。Faster Transformer 对外提供 C++ API，TensorFlow OP 接口，以及 TensorRT [4] 插件，并提供了相应的示例，用以支持用户将其集成到不同的线上应用代码中。

Faster Transformer 目前已经开源，可以访问https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer

获取项目全部源代码，最新的性能数据以及支持的特性。欢迎大家前往使用，加星和反馈。

性能数据

Faster Transformer 在不同的应用场景下都有着突出的表现。我们在这里测试了不同生产环境下 Faster Transformer 前向计算的执行时间以及与 TensorFlow XLA 的性能比较。测试环境如表 1 所示：

表1. 性能数据测试环境（本地服务器）

首先针对线上 QPS 较低的业务（例如问答），我们将 batch size 设置为 1，测试了 BERT 标准模型在不同的句子长度下，12 层 Transformer 在 P4 和 T4 上的性能。由于这种场景下 TensorFlow 的性能非常依赖于 CPU，因此这里不予列出。

表2. 小 batch size 情况下 Faster Transformer 的性能

接着我们来观察 Faster Transformer 在搜索或者广告推荐等大 batch size 场景下的加速效果。表 3 和表 4分别测试了固定句子长度为 32，标准模型（12 head x 64维）和裁剪模型（4 head x 32维）在不同 batch size下，12 层 Transformer 在 V100 上使用了 FP16 计算精度的性能。

表3. 标准模型不同 Batch Size下 TensorFlow XLA 和 Faster Transformer 在 V100 上的性能对比

表4. 裁剪模型不同 Batch Size下TensorFlow XLA 和 Faster Transformer 在 V100 上的性能对比

可以看出，在标准模型和裁剪模型上，Faster Transformer 都有很好的加速效果。

使用方法

Faster Transformer 提供了 TensorFlow OP ，C++ API 和 TensorRT Plugin 三种接口。

在 TensorFlow 中使用 Faster Transformer

在 TensorFlow 中使用 Faster Transformer 最为简单。只需要先 import .so 文件，然后在代码段中添加对 Faster Transformer OP 的调用即可。具体代码如下所示。

# import op
transformer_op_module = tf.load_op_library(os.path.join('../../build/lib/libtf_transformer.so'))
...
def fast_transformer_model_trans(...)
    ...
      # original code
      ...
      layer_output = layer_norm(layer_output + attention_output)

      # calling faster transformer op
      trainable_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=tf.get_variable_scope().name)
      layer_output = transformer_op_module.bert_transformer(
        layer_input,
        layer_input,
        trainable_vars[0], trainable_vars[2], trainable_vars[4], trainable_vars[1], trainable_vars[3], trainable_vars[5], 
        attention_mask,
        trainable_vars[6], trainable_vars[7], trainable_vars[8], trainable_vars[9], trainable_vars[10], trainable_vars[11],
        trainable_vars[12], trainable_vars[13], trainable_vars[14], trainable_vars[15],
        batch_size=batch_size, from_seq_len=seq_length, to_seq_len=seq_length, head_num=num_attention_heads, size_per_head=attention_head_size)
      
      # original code
      ...
      all_layer_outputs.append(layer_output)
    ...

使用 C++ API 或者 TensorRT 调用 Faster Transformer

考虑到封装成 TensorFlow OP 会引入一些额外的开销，我们更建议用户直接使用 C++ API 或者 TensorRT Plugin 的方式去集成。目前这两种方式不支持直接解析训练好的模型。Transformer 层所需要的 weights 参数，需要用户手动从训练好的模型中导出。调用方式相对简单，将导出的 weights 赋值给参数结构体，创建相应的对象，调用 initialize 或者 build_engine 函数初始化对象。运行时，每次调用 forward 或者do_inference 即可。具体代码如下所示。

/* C++ interface */
typedef BertEncoderTransformerTraits<OperationType::HALF,  cuda::OpenMultiHeadAttention> EncoderTraits_;
fastertransformer::Allocator<AllocatorType::CUDA> allocator(0);
EncoderInitParam<__half> encoder_param; //init param here

encoder_param.from_tensor = d_from_tensor;
...

BertEncoderTransformer<EncoderTraits_> *encoder_transformer_ = new 
 BertEncoderTransformer<EncoderTraits_>(allocator, batch_size, from_seq_len, to_seq_len, head_num, size_per_head);
encoder_transformer_->initialize(encoder_param);
encoder_transformer_->forward();

/* TensorRT Plugin */
std::vector<std::vector<T *> > params;
/* push all weight pointers into the vector params*/
TRT_Transformer<T>* trt_transformer = new TRT_Transformer<T>(batch_size, seq_len, head_num, hidden_dim, layers);
trt_transformer->build_engine(params);
trt_transformer->do_inference(batch_size, h_from_tensor, h_attr_mask, h_transformer_out, stream);

优化原理

在深入了解 Faster Transformer 的优化原理之前，我们先来看下 TensorFlow 的实现情况。图1是 TensorFlow 在默认计算模式（不使用 XLA 优化）下的时间线片段。

图1. TensorFlow计算GELU的时间线

其中，黄色矩形框中对应的是激活函数 GELU。可以看到，在 TensorFlow 中，这个函数是通过 8 个类似Pow，Add，和 Tanh 等基本 OP 来实现的。Layer Normalization 操作也是类似的情况（图2）。

图2. TensorFlow计算Layer Normalization的时间线

在 TensorFlow 中，每一个基本 OP 都会对应一次 GPU kernel 的调用，和多次显存读写，这些都会增加大量额外的开销。TensorFlow XLA 可以在一定程度上缓解这个问题，它会对一些基本的 OP 进行合并，以减少GPU kernel 的调度和显存读写。但在大多数情况下，XLA 依然无法达到最优的性能，特别是对于 BERT 这种计算密集的情况，任何性能的提升都将节省巨量的计算资源。

如我们前面提到的，OP 融合可以降低 GPU 调度和显存读写，进而提升性能。出于性能最大化的考虑，在Faster Transformer 内部，我们将除矩阵乘法以外的所有 kernel 都进行了尽可能的融合，单层Transformer 的计算流程如下图所示：

图3. BERT中Transformer Layer 的计算流程图

如图3所示，Faster Transformer 只用了 14 个 kernel 就完成了原来将近 60 个 kernel 的计算逻辑。这其中，8 个 kernel 是通过调用 cuBLAS 接口计算矩阵乘法（绿色框），其余 6 个是自定义 kernel （蓝色框）。

针对 batch size 比较小的场景（例如问答，TTS 等），简单的融合后，基本上就可以达到很好的性能。这类场景下，TensorFlow 原生实现的最大瓶颈就在于频繁的 kernel launch，融合后大大降低了 launch 的开销，因此可以比较轻易地获得很好的加速效果。

针对大 batch 的场景，我们需要对矩阵乘法和所有的自定义 kernel 做精细的调优，才能达到很好的加速效果。我们从矩阵乘法算法选择，非矩阵乘法操作的参数配置，SoftMax 多版本实现，以及数据结构类型等几个方面对大 batch 的情况进行了专门的调优。

首先针对矩阵乘法，在调用 cuBLAS 的接口时，可以指定性能最优的算法。特别是针对 Volta 和 Turing 架构的 GPU，使用 Tensor Core 进行半精度计算时，当精度满足需求的情况下，累加器也可以选择半精度，从而进一步提升性能。

除矩阵乘法以外的 6 个 kernel，大部分都是对矩阵乘的结果进行一些 element-wise 的操作。输入矩阵的大小，跟 4 个参数有关，batch size，句子长度，attention 的 head 数量以及每个 head 的维度。针对不同的应用场景，参数大小可能极为不同。比如在线问答类的场景，batch size 可能为会很小，通常为 1。而广告推荐或者搜索类的场景，batch size 通常跟候选集大小有关，一般会是几百的规模。这样，输入矩阵的行数变化范围可能是几十到上千。因此，我们需要针对不同的情况，动态的调整 kernel launch 时的配置参数（grid 和 block 的大小），甚至要针对同一个功能实现多个不同版本的 kernel 函数，例如，SoftMax 的计算就有两个不同的实现版本。

针对半精度 FP16，我们对各个 kernel 也进行了相应优化。首先，在 kernel 的实现中，将输入的 half 指针转成 half2 类型，并使用了 half2 相关的数学函数。这样不仅仅可以达到 2 倍于 half 的访存带宽和计算吞吐，还可以极大地减少指令的发射数量。其次，在 SoftMax 以及 Layer Normalization 的操作中，为防止求和溢出，将数据以 half2 的形式读入后，会转成 float2 类型，来做求和计算。

除上述优化之外，Faster Transformer 还优化了前向计算中耗时较高的 GELU 激活函数，Layer Normalization 以及 SoftMax 等操作。比如利用 warp shuffle 实现高效的矩阵按行求和操作， 将 1/sqrtf计算替换为 rsqrtf 函数，以及 power (x, 3.0) 替换为x * x * x等。总之，我们针对 Transformer 进行了各种优化以保证它的高效执行。

结论

Faster Transformer 是一个开源的高效 Transformer 实现，相比 TensorFlow XLA  可以带来 1.5-2x 的提速。Faster Transformer 对外提供 C++ API， TensorFlow OP，以及 TensorRT Plugin 三种接口。对每种接口的调用方式，我们提供了完整的示例，方便用户集成。

Faster Transformer 目前已经开源，可以访问https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer

获取项目全部源代码，最新的性能数据以及支持的特性。欢迎大家前往使用，加星和反馈。

[1] Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. “Attention Is All You Need.” ArXiv:1706.03762 [Cs], June 12, 2017. http://arxiv.org/abs/1706.03762.

[2] Devlin, Jacob, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. “BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding.” ArXiv:1810.04805 [Cs], October 10, 2018. http://arxiv.org/abs/1810.04805.

[3] Yang, Zhilin, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. “XLNet: Generalized Autoregressive Pretraining for Language Understanding.” ArXiv:1906.08237 [Cs], June 19, 2019. http://arxiv.org/abs/1906.08237.

[4] TensorRT: https://developer.nvidia.com/tensorrt

[5] Turing T4 GPU, more information: https://www.nvidia.com/en-us/data-center/tesla-t4/

[6] Volta V100 GPU, more information: https://www.nvidia.com/en-us/data-center/tesla-v100/

[7] Pascal P4 GPU, more information: https://www.nvidia.com/en-us/deep-learning-ai/solutions/inference-platform/hpc/

（本文作者：NVIDIA GPU计算专家团队，贾晓莹）

关于NVIDIA

NVIDIA（纳斯达克股票代码：NVDA）在1999年发明的GPU激发了PC游戏市场的增长，重新定义了现代计算机显卡，并且对并行计算进行了革新。最近，通过将GPU作为可以感知和理解世界的计算机、机器人乃至自动驾驶汽车的大脑，GPU深度学习再度点燃了全新的计算时代——现代人工智能。更多信息，请访问http://nvidianews.nvidia.com/。