关于开放源代码:速度超快字节跳动开源序列推理引擎LightSeq

这应该是业界第一款残缺反对 Transformer、GPT 等多种模型高速推理的开源引擎。

2017 年 Google 提出了 Transformer [1] 模型，之后在它根底上诞生了许多优良的预训练语言模型和机器翻译模型，如 BERT [2]、GPT 系列 [13] 等，一直刷新着泛滥自然语言解决工作的能力程度。与此同时，这些模型的参数量也在出现近乎指数增长（如下图所示）。例如最近引发热烈探讨的 GPT-3 [3]，领有 1750 亿参数，再次刷新了参数量的记录。

如此微小的参数量，也为模型推理部署带来了挑战。以机器翻译为例，目前 WMT[4]较量中 SOTA 模型曾经达到了 50 层以上。支流深度学习框架下，翻译一句话须要好几秒。这带来了两个问题：一是翻译工夫太长，影响产品用户体验；二是单卡 QPS（每秒查问率）太低，导致服务老本过高。

因而，明天给大家安利一款速度十分快，同时反对十分多个性的高性能序列推理引擎——LightSeq。它对以 Transformer 为根底的序列特征提取器（Encoder）和自回归的序列解码器（Decoder）做了深度优化，早在 2019 年 12 月就曾经开源，利用在了包含火山翻译等泛滥业务和场景。据理解，这应该是业界第一款残缺反对 Transformer、GPT 等多种模型高速推理的开源引擎。

LightSeq 能够利用于机器翻译、主动问答、智能写作、对话回复生成等泛滥文本生成场景，大大提高线上模型推理速度，改善用户的应用体验，升高企业的经营服务老本。

相比于目前其余开源序列推理引擎，LightSeq 具备如下几点劣势：

1. 高性能

LightSeq 推理速度十分快。例如在翻译工作上，LightSeq 相比于 Tensorflow 实现最多能够达到 14 倍的减速。同时当先目前其余开源序列推理引擎，例如最多可比 Faster Transformer 快 1.4 倍。

1. 反对模型性能多

LightSeq 反对 BERT、GPT、Transformer、VAE 等泛滥模型，同时反对 beam search、diverse beam search[5]、sampling 等多种解码形式。下表具体列举了 Faster Transformer[7]、Turbo Transformers[6]和 LightSeq 三种推理引擎在文本生成场景的性能差别：

1. 简略易用，无缝连接 Tensorflow、PyTorch 等深度学习框架

LightSeq 通过定义模型协定，反对各种深度学习框架训练好的模型灵便导入。同时蕴含了开箱即用的端到端模型服务，即在不须要写一行代码的状况下部署高速模型推理，同时也灵便反对多层次复用。

应用办法

利用 LightSeq 部署线上服务比拟简便。LightSeq 反对了 Triton Inference Server[8]，这是 Nvidia 开源的一款 GPU 推理 server，蕴含泛滥实用的服务中间件。LightSeq 反对了该 server 的自定义推理引擎 API。因而只有将训练好的模型导出到 LightSeq 定义的模型协定 [9] 中，就能够在不写代码的状况下，一键启动端到端的高效模型服务。更改模型配置（例如层数和 embedding 大小）都能够不便反对。具体过程如下：

首先筹备好模型仓库，上面是目录构造示例，其中 transformer.pb 是按模型协定导出的模型权重，libtransformer.so 是 LightSeq 的编译产物。

 - model_zoo/
 - model_repo/
 - config.pbtxt
 - transformer.pb
 - 1/
 - libtransformer.so

而后就能够启动 Triton Inference Server[8]，搭建起模型服务。

1. trtserver --model-store=${model_zoo}

性能测试

在 NVIDIA Tesla P4 和 NVIDIA Tesla T4 显卡上，笔者测试了 LightSeq 的性能，抉择了深度学习框架 Tensorflow v1.13 和解码场景反对较为丰盛的 Faster Transformer v2.1 实现作为比照。Turbo Transformers 解码办法比拟繁多（只反对 Beam Search，不反对文本生成中罕用的采样解码），尚未满足理论利用需要，因而未作比照。

机器翻译性能

在机器翻译场景下，笔者测试了 Transformer base 模型（6 层 encoder、6 层 decoder、隐层维度 512）采纳 beam search 解码的性能，试验后果如下：

能够发现，在小 batch 场景下，Faster Transformer 和 LightSeq 比照 Tensorflow 都达到了 10 倍左右的减速。而随着 batch 的增大，因为矩阵乘法运算占比越来越高，两者对 Tensorflow 的减速比都呈衰减趋势。LightSeq 衰减绝对平缓，特地是在大 batch 场景下更加具备劣势，最多能比 Faster Transformer 快 1.4 倍。这也对将来的一些推理优化工作提供了领导：小 batch 场景下，只有做好非计算密集型算子交融，就能够获得很高的减速收益；而大 batch 场景下则须要持续优化计算密集型算子，例如矩阵乘法等。

最初在 WMT14 规范的法英翻译工作上，笔者测试了 Transformer big 模型的性能。LightSeq 在 Tesla P4 显卡上均匀每句翻译提早为 167ms，Tesla T4 上减小到了 82ms。而作为比照，TensorFlow 提早均为 1071ms，LightSeq 别离达到了 6.41 和 13.06 倍减速。另外，笔者尝试了其余多种模型配置，失去了比拟统一的减速效率。例如更深层的模型构造上（encoder 加深至 16 层），LightSeq 失去的减速比，别离是 6.97 和 13.85 倍。

文本生成性能

上述机器翻译通常采纳 Beam Search 办法来解码，而在文本生成场景，常常须要应用采样（Sampling）来晋升生成后果的多样性。下图展现了 Transformer base 模型采纳 top-k/top-p sampling 的性能测试比照：

能够发现，在须要应用采样解码的工作中，LightSeq 在大部分配置下当先于 Faster Transformer，最多也能达到 1.4 倍的额定减速。此外，相比于 TensorFlow 实现，LightSeq 对 GPT 和 VAE 等生成模型也达到了 5 倍以上的减速成果。

服务压力测试

在云服务上，笔者测试了在理论利用中 GPT 场景下，模型服务从 Tensorflow 切换到 LightSeq 的提早变动状况（服务显卡应用 NVIDIA Tesla P4）。能够察看到，pct99 提早升高了 3 到 5 倍，峰值从 360 毫秒左右降落到 80 毫秒左右，具体后果如下图所示：

更多的比照试验后果能够在 LightSeq 性能评测报告 [10] 中查看到。

技术原理

以 Transformer 为例，一个机器翻译 / 文本生成模型推理过程包含两局部：序列编码模块特色计算和自回归的解码算法。其中特色计算局部以自注意力机制及特色变换为外围（矩阵乘法，计算密集型），并随同大量 Elementwise（如 Reshape）和 Reduce（如 Layer Normalization）等 IO 密集型运算；解码算法局部蕴含了词表 Softmax、beam 筛选、缓存刷新等过程，运算琐碎，并引入了更简单的动静 shape。这为模型推理带来了泛滥挑战：

1.  IO 密集型计算的细粒度核函数调用带来大量冗余显存读写，成为特色计算性能瓶颈。
2. 简单动静 shape 为计算图优化带来挑战，导致模型推理期间大量显存动静申请，耗时较高。
3. 解码生成每一步字符过程逻辑简单，难以并行化计算从而施展硬件劣势。

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 首先对模型中所有动静的 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

topk_log_probs, topk_indices = tf.nn.top_k(log_seq_prob, k=K)

3. 依据 beam id，刷新 decoder 中的 self attention 模块中的 key 和 value 的缓存

refresh_cache(cache, topk_indices)

能够发现，为了筛选概率 top-k 的 token，必须在 [batch_size, beam_size, vocab_size]大小的 logit 矩阵上进行 softmax 计算及显存读写，而后进行 batch_size 次排序。通常 vocab_size 都是在几万规模，因而计算量十分宏大，而且这仅仅只是一步解码的计算耗费。因而实际中也能够发现，解码模块在自回归序列生成工作中，累计提早占比很高（超过 30%）。

LightSeq 的翻新点在于联合 GPU 计算个性，借鉴搜寻举荐中罕用的粗选 - 精排的两段式策略，将解码计算改写成层级式，设计了一个 logit 粗选核函数，胜利防止了 softmax 的计算及对十几万元素的排序。该粗选核函数遍历 logit 矩阵两次：

• 第一次遍历，对每个 beam，将其 logit 值随机分成 k 组，每组求最大值，而后对这 k 个最大值求一个最小值，作为一个近似的 top- k 值（肯定小于等于实在 top- k 值），记为 R -top-k。在遍历过程中，同时能够计算该 beam 中 logit 的 log_sum_exp 值。

• 第二次遍历，对每个 beam，找出所有大于等于 R-top-k  的 logit 值，将（logit – log_sum_exp + batch_id offset, beam_id vocab_size + vocab_id）写入候选队列，其中 offset 是 logit 的下界。

在第一次遍历中，logit 值通常遵从正态分布，因而算出的 R -top- k 值十分靠近实在 top- k 值。同时因为这一步只波及到寄存器的读写，且算法复杂度低，因而能够疾速执行实现（十几个指令周期）。理论察看发现，在 top- 4 设置下，依据 R -top- k 只会从几万 token 中粗选出十几个候选，因而十分高效。第二次遍历中，依据 R -top- k 粗选出候选，同时对 logit 值按 batch_id 做了值偏移，多线程并发写入显存中的候选队列。

粗选实现后，在候选队列中进行一次排序，就能失去整个 batch 中每个序列的精确 top- k 值，而后更新缓存，一步解码过程就疾速执行实现了。

上面是 k =2，词表大小 = 8 的状况下一个具体的示例（列代表第几个字符输入，行代表每个地位的候选）。能够看出，原来须要对 16 个元素进行排序，而采纳层级解码之后，最初只须要对 5 个元素排序即可，大大降低了排序的复杂度。

可视化剖析计算提早

为了验证下面几种优化技术的实际效果，笔者用 GPU profile 工具，对 LightSeq 的一次推理过程进行了提早剖析。下图展现了 32 位浮点数和 16 位浮点数精度下，各计算模块的提早占比：

能够发现，在两种计算精度下：

1. 通过优化后，cuBLAS  中的矩阵乘法计算提早别离占比 82% 和 88%，成为推理减速新的次要瓶颈。而作为比照，咱们测试了 Tensorflow 模型，矩阵乘法计算提早只占了 25%。这阐明 LightSeq 的 beam search 优化曾经将提早降到了非常低的程度。
2. 缓存刷新别离占比 10% 和 6%，比重也较高，但很难持续优化。今后能够尝试缩小缓存量（如升高 decoder 层数，升高缓存精度等）来持续升高提早。
3. 其余运算总计占比 8% 和 6%，包含了 Layer Normalization、beam search 和两头后果的显存读写等。

可视化后果阐明了 LightSeq 曾经做到了极致优化，大大晋升了推理速度。

传送门：

GitHub 我的项目地址：

https://github.com/bytedance/…

[1] Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems. 2017.

[2] Devlin, Jacob, et al. “Bert: Pre-training of deep bidirectional transformers for language understanding.” arXiv preprint arXiv:1810.04805 (2018).

[3] Brown, Tom B., et al. “Language models are few-shot learners.” arXiv preprint arXiv:2005.14165 (2020).

[4] WMT2020, http://www.statmt.org/wmt20/

[5] Li, Jiwei, Will Monroe, and Dan Jurafsky. “A simple, fast diverse decoding algorithm for neural generation.” arXiv preprint arXiv:1611.08562 (2016).

[6] TurboTransformers, https://github.com/Tencent/Tu…

[7] FasterTransformer, https://github.com/NVIDIA/Dee…

[8] NVIDIA Triton Inference Server, https://github.com/triton-inf…

[9] LightSeq proto, https://github.com/bytedance/…

[10] LightSeq 性能评测报告, https://github.com/bytedance/…

[11] LightSeq Layer Normalization, https://github.com/bytedance/…

[12] cuBLAS, https://docs.nvidia.com/cuda/…

[13] GPT2,”Language Models are Unsupervised Multitask Learners”