Transformer离线部署-GPU优化策略

前言<a id="sec-1"></a>

本文主要介绍Transformer类网络在GPU设备上部署上的优化要点。 主要围绕Nvidia开源的 FasterTransformer展开。

模型结构分析<a id="sec-2"></a>

标准的Transformer结构主要包括 Encoder 和 Decoder 两部分结构，具体结构分析可参考 Transformer在CV领域的应用与部署：

Encoder对应算子结构为：

Decoder对应算子结构为：

可以发现：

Encoder/Decoder的基本结构还是一致的。其中主要算子（算力大头）为GEMM/Softmax算子；另外存在大量add bias/normalization操作（包含大量小算子）。从 Encoder 模型下面流水图中可以发现：大量的小算子会导致CPU侧的算子"Launch"时间变长，呈现“Launch Bound”的现象，无法完全发挥GPU的算力。

具体优化措施<a id="sec-3"></a>

基于Transformer类网络的基本结构，当前NVIDIA做了一系列的优化。首先是算子融合层面，可以进行一系列的小算子融合：

1. LayerNorm小算子融合：add bias + layernorm
2. 激活函数小算子融合：add bias + relu/gelu
3. Softmax小算子融合：add bias + Softmax

另外核心算子层面：

1. GEMM为主要耗时算子，cuBLAS着力进行了性能优化：QKV GEMM融合等/GEMM配置自动选择等
2. 采用了优化的Softmax算子：FP32/FP16实现不同

其他优化点：

1. 采用了half2数据类型，能够减少数据拷贝的空间
2. 指令层面：
   
1. _expf替换expf运算：精度更低，但是性能更好（主要提速softmax？）
2. 采用 __shfl_xor_sync 对 reduce 操作进行加速（Layernorm）

优化后的 Encoder 结构流水图如下，可见各算子之间基本不存在空隙，提高了GPU的利用率；同时对核心算子的优化也大大提升了模型性能。

参考资料<a id="sec-4"></a>

FasterTransformer Presentation-YouTube