【NLP】Fastertransformer源码解读

最近拜读了NVIDIA前阵子开源的fastertransformer，对CUDA编程不是很熟悉，但总算是啃下来一些，带你们读一下硬核源码。

1. 简介

英伟达公众号推送的文章加上配图其实已经把该要讲的很清楚了，主要有如下几方面：

1. 为了减小kernel调用次数，将除了矩阵乘法的kernel都尽量合并
2. 针对大batch单独进行了kernel优化
3. 支持选择最优的矩阵乘法
4. 在使用FP16时使用half2类型，达到half两倍的访存带宽和计算吞吐
5. 优化gelu、softmax、layernorm的实现以及选用rsqrt等

不了解底层的同窗可能不是很懂，没事我刚看到的时候也不懂，也不敢问，强撸一下源码就通透（fang qi）了

2. 硬核源码解读

首先简略说一下第一点优化。Kernel在tensorflow里的概念是operation的计算实现，在cuda里是执行一个线程的函数，也是一次计算，只不过tensorflow的更加宏观些。每次tensorflow执行一个operation，都要调用对应的OpKernel，试想一个经过TF实现的transformer，有将近60个operation，计算一次要执行60次上述过程，进行频繁的GPU调度和显存读写。所以fastertransformer尽量多地对kernel进行了合并。

2.1 总体结构

Fastertransformer目录下主要有（如下简称FTF）：

1. fastertransformer：主要源码
1. cuda：优化后的kernel以及对multi-head attention总体的封装（没过线性层）
2. tf_op：tensorflow operation和OpKernel的注册（op理解为声明、Opkenerl是定义）
3. trt_plugin：tensorRT的实现（能够支持multi streaming太赞了）
4. bertencodertransformer.h：transformer总体的封装
2. sample：cpp、tensorflow、tensrflow_bert、tensorRT的调用FTF的示例
3. tools：根据参数选择最优的矩阵乘法（GEMM=General Matrix Multiplication）

接下来我主要想讲一下1.1的细节，1.2能够参考这位大佬的文章，剩下的代码可读性很强，基本读一两遍就知道了。

除去矩阵乘法，做者把剩下的op合并成了4个（图中蓝色框）：

这四个op的cuda源码分别在cuda_kernels.cu和open_attention.cu两个文件中，接下来研究下每一个op。

2.2 add_QKVbias (open_attention.cu)

在FP32时，每一个block负责处理一个word(num_head*head_size)的add bias运算，先找到要处理QKV中的一个，再进行运算，由于要transpose，因此把结果存入[bsz, num_head, seq_len, head_size]的矩阵里。

在FP16是每一个block同时处理QKV上的同一个word（多是由于fp16计算的更快一些），在实际的计算中把half都转成了half2计算。add的话直接用封装好的__hadd2运算。使用half2计算的缘由原文说的比较清楚：

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

2.3 softmax_kernel (open_attention.cu)

在计算softmax以前对block线程数进行了区间处理，由于block里的线程数最好是wrap大小(32)的倍数，提升计算效率。

调用kernel以前，会根据batch_size * head_num选择不一样的softmax kernel，主要是为了保证在大batch的状况下的计算效率，这里为何使用120我也不是很清楚，但愿懂的朋友助力一下

if(batch_size * head_num <= 120)
    {
      grid.x = batch_size * head_num * seq_len;
      softmax_kernel_v2<DataType_><<<grid, block, 0, stream>>>(qk_buf_, attr_mask, batch_size, head_num, seq_len, scaler); 
    }
    else
    {
      grid.x = batch_size * head_num;
      softmax_kernel<DataType_><<<grid, block, 0, stream>>>(qk_buf_, attr_mask, batch_size, head_num, seq_len, scaler); 
    }复制代码

在算softmax时，分母有个求和操做，用到了经典的parallel reduce算法，能够仔细看看参考，讲的比较清楚。

这里注意，使用最第一版源码的同窗们须要照着实现一个blockReduceMax，以防止数值溢出，softmax严谨的实现应该是：

def softmax(x):
"""Compute the softmax of vector x."""
    exp_x = np.exp(x)
    softmax_x = exp_x / np.sum(exp_x)
    return softmax_x复制代码

2.4 transpose (open_attention.cu)

这里要transpose回[bsz, seq_len, num_head, head_size]的矩阵。由于c++里面矩阵是行优先存储，只要按顺序乘过来就行了（最开始看的有点晕）。

2.5 add_bias_act & add_bias_input_layernorm (cuda_kernels.cu)

若是前面几个函数啃下来了，这两个就比较好懂，主要的优化点是：

1. x^3 -> x*x*x: c语言中x*x和pow(x,2)哪一个计算更快一点？
2. rsqrt: Why is SSE scalar sqrt(x) slower than rsqrt(x) * x?
3. 还有就是各类half2运算的使用

2.6 gemm (tools/gemm_test)

矩阵乘法根据fp16和fp32的不一样在不一样的cublas算法中选择，选择后记录到http://gemm_config.in文件中：

问了下做者，其实fp32也可使用CUBLAS_GEMM_ALGO0_TENSOR_OP到CUBLAS_GEMM_ALGO15_TENSOR_OP的算法，只不过存在一些风险（使用后速度提高2倍）。

2.7 trt_plugin

做者额外封装了一个tensorRT的层，tensorRT主要是经过engine，在给定的context和stream下进行异步计算，提供了multi stream inference的可能。关于TensorRT的异步编程推荐一个英伟达的PPT：

CUDA C/C++ Streams and Concurrency​ developer.download.nvidia.cn

---

这篇文章写做周期比较长，主要是源码比较硬核，边看边学cuda和c++，到如今也就懂了80%左右吧。不过fastertransformer是真的香，并且直接用tensorflow也很方便，各位须要inference的朋友们必定要用呀