这应该是业界第一款残缺反对 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"