Press "Enter" to skip to content

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

 

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. 性能数据测试环境(本地服务器)

 

软件版本CUDA 10.0
TensorFlow 1.13
硬件参数CPU: Intel(R) Xeon(R) Gold 6132 CPU @ 2.60GHz
Turing T4[5] @mclk 5000MHz, pclk 1590MHz

Volta V100[6] @ mclk 877MHz, pclk 1380MHz

Pascal P4[7] @ mclk 2999MHz, pclk 1531MHz

 

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

 

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

 

 

batch size = 1, number of heads = 12, size per head = 64, 12 layers, time in ms
Sequence LengthP4 in FP32T4 in FP32T4 in FP16
323.42.71.6
644.03.61.8
1286.25.92.2

 

接着我们来观察 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 上的性能对比

 

 

Sequence length = 32, number of heads = 12, size per head = 64, 12 layers
Batch sizeTensorFlow XLA (ms)Faster Transformer (ms)Speedup
10014.09.61.5x
20026.518.41.5x
30038.427.41.5x
40049.735.61.5x
50062.244.61.5x

 

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

 

 

Sequence length = 32, number of heads = 4, size per head = 32, 12 layers
Batch sizeTensorFlow XLA (ms)Faster Transformer (ms)Speedup
1003.51.72.0x
2004.92.61.9x
006.43.41.9x
4008.04.31.9x
5009.95.11.9x

 

可以看出,在标准模型和裁剪模型上,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/

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注