【NLP】Fastertransformer源码解读

1,783 阅读5分钟

最近拜读了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 Concurrencydeveloper.download.nvidia.cn

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