江苏龙网

  1. 首页 > 资讯 > 科技 > >

|速度超快！字节跳动开源序列推理引擎LightSeq( 三 )


LightSeq 取得这么好的推理加速效果 ， 对这些挑战做了哪些针对性的优化呢？笔者分析发现 ， 核心技术包括这几项：融合了多个运算操作来减少 IO 开销、复用显存来避免动态申请、解码算法进行层级式改写来提升推理速度 。 下面详细介绍下各部分的优化挑战和 LightSeq 的解决方法 。
算子多运算融合
近年来 ， 由于其高效的特征提取能力 ， Transformer encoder/decoder 结构被广泛应用于各种 NLP 任务中 ， 例如海量无标注文本的预训练 。 而多数深度学习框架（例如 Tensorflow、Pytorch 等）通常都是调用基础运算库中的核函数（kernel function）来实现 encoder/decoder 计算过程 。 这些核函数往往粒度较细 ， 通常一个组件需要调用多个核函数来实现 。
以层归一化（Layer Normalization）为例 ， Tensorflow 是这样实现的：
mean = tf.reduce_mean(x, axis=[-1], keepdims=True)variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)result = (x - mean) * tf.rsqrt(variance + epsilon) * scale + bias
可以发现 ， 即使基于编译优化技术（自动融合广播（Broadcast）操作和按元素（Elementwise）运算） ， 也依然需要进行三次核函数调用（两次 reduce_mean ， 一次计算最终结果）和两次中间结果的显存读写（mean 和 variance） 。 而基于 CUDA ， 我们可以定制化一个层归一化专用的核函数 ， 将两次中间结果的写入寄存器 。 从而实现一次核函数调用 ， 同时没有中间结果显存读写 ， 因此大大节省了计算开销 。 有兴趣的同学可以在文末参考链接中进一步查看具体实现[11] 。
基于这个思路 ， LightSeq 利用 CUDA 矩阵运算库 cuBLAS[12]提供的矩阵乘法和自定义核函数实现了 Transformer ， 具体结构如下图所示：
|速度超快！字节跳动开源序列推理引擎LightSeq
本文插图

蓝色部分是自定义核函数 ， 黄色部分是矩阵乘法 。 可以发现 ， 矩阵乘法之间的运算全部都用一个定制化核函数实现了 ， 因此大大减少了核函数调用和显存读写 ， 最终提升了运算速度 。
动态显存复用
为了避免计算过程中的显存申请释放并节省显存占用 ， LightSeq 首先对模型中所有动态的 shape 都定义了最大值（例如最大序列长度） ， 将所有动态shape转换为静态 。 接着在服务启动的时候 ， 为计算过程中的每个中间计算结果按最大值分配显存 ， 并对没有依赖的中间结果共用显存 。 这样对每个请求 ， 模型推理时不再申请显存 ， 做到了：不同请求的相同 Tensor 复用显存；同请求的不同 Tensor 按 shape 及依赖关系复用显存 。
通过该显存复用策略 ， 在一张 T4 显卡上 ， LightSeq 可以同时部署多达 8 个 Transformer big 模型（batch_size=8 ， 最大序列长度=8 ， beam_size=4 ， vocab_size=3万） 。 从而在低频或错峰等场景下 ， 大大提升显卡利用率 。
层级式解码计算
在自回归序列生成场景中 ， 最复杂且耗时的部分就是解码 。 LightSeq 目前已经支持了 beam search、diversity beam search、top-k/top-p sampling 等多种解码方法 ， 并且可以配合 Transformer、GPT使用 ， 达到数倍加速 。 这里我们以应用最多的 beam search 为例 ， 介绍一下 LightSeq 对解码过程的优化 。
首先来看下在深度学习框架中传统是如何进行一步解码计算的：
# 1.计算以每个token为结尾的序列的log probability
log_token_prob = tf.nn.log_softmax(logit) # [batch_size, beam_size, vocab_size]log_seq_prob += log_token_prob # [batch_size, beam_size, vocab_size]log_seq_prob = tf.reshape(log_seq_prob, [-1, beam_size * vocab_size])
# 2. 为每个序列（batch element）找出排名topk的token


推荐阅读


上一篇：|人工智能小冰：从音乐创作者到“音乐科技领军人物”

下一篇：极果|太鸡贼！非iPhone12手机MageSafe充电极慢，iPhone11也不例外