正文

字节跳动开源序列推理引擎lightseq(代码片段)

Alex_996  Alex_996  2022-12-03  543

关键词：

背景介绍

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具有如下几点优势：

高性能。
LightSeq推理速度非常快。例如在翻译任务上，LightSeq相比于Tensorflow实现最多可以达到14倍的加速。同时领先目前其他开源序列推理引擎，例如最多可比Faster Transformer快1.4倍。
支持模型功能多。
LightSeq支持BERT、GPT、Transformer、VAE 等众多模型，同时支持beam search、diverse beam search[5]、sampling等多种解码方式。下表详细列举了Faster Transformer[7]、Turbo Transformers[6]和LightSeq三种推理引擎在文本生成场景的功能差异：
简单易用，无缝衔接Tensorflow、PyTorch等深度学习框架。
LightSeq通过定义模型协议，支持各种深度学习框架训练好的模型灵活导入。同时包含了开箱即用的端到端模型服务，即在不需要写一行代码的情况下部署高速模型推理，同时也灵活支持多层次复用。

使用方法

使用python调用LightSeq

我们提供一个huggingface bart的例子，这样可以比较Lightseq在序列生成上的优异性能。首先，可以安装下列依赖

pip install torch tensorflow transformers lightseq
cd example/python

接着运行下面两个脚本就可以看到LightSeq和Huggingface的性能对比。hf_bart_export.py 用来将huggingface的模型参数导出成protobuffer文件供LightSeq使用。

python hf_bart_export.py
python ls_bart.py

在我们的Tesla V100显卡上能获得下面的输出，可以看到LightSeq对比huggingface有47倍左右的加速比。

=========================lightseq=========================
lightseq generating...
lightseq time: 0.034502994269132614s
lightseq results:
I love that girl, but she does not love me.
She is so beautiful that I can not help glance at her.
Nothing's gonna change my love for you.
Drop everything now. Meet me in the pouring rain. Kiss me on the sidewalk.
=========================huggingface=========================
huggingface generating...
huggingface time: 1.6297104470431805s
huggingface results:
I love that girl, but she does not love me.
She is so beautiful that I can not help glance at her.
Nothing's gonna change my love for you.
Drop everything now. Meet me in the pouring rain. Kiss me on the sidewalk.

目前从pypi安装LightSeq只支持Linux 3.6-3.8的python版本，后面我们会尝试支持更多的平台。

使用NV Triton inference server部署

利用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]，搭建起模型服务。

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。这为模型推理带来了众多挑战：

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

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位浮点数精度下，各计算模块的延迟占比：

可以发现，在两种计算精度下：

经过优化后，cuBLAS中的矩阵乘法计算延迟分别占比82%和88%，成为推理加速新的主要瓶颈。而作为对比，我们测试了Tensorflow模型，矩阵乘法计算延迟只占了25%。这说明LightSeq的beam search优化已经将延迟降到了非常低的水平。
缓存刷新分别占比10%和6%，比重也较高，但很难继续优化。今后可以尝试减少缓存量（如降低decoder层数，降低缓存精度等）来继续降低延迟。
其他运算总计占比8%和6%，包括了Layer Normalization、beam search和中间结果的显存读写等。

可视化结果说明了LightSeq已经做到了极致优化，大大提升了推理速度。

传送门

GitHub项目地址：https://github.com/bytedance/lightseq

[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/TurboTransformers
[7] FasterTransformer, https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer
[8] NVIDIA Triton Inference Server, https://github.com/triton-inference-server/server
[9] LightSeq proto, https://github.com/bytedance/lightseq/blob/79249ada7f7078b685423d9a19f85e454217d5a4/lightseq/inference/proto/
[10] LightSeq性能评测报告, https://github.com/bytedance/lightseq/blob/f713ba960985b572c477a8279019b80344eb1c7f/docs/inference/performance.md
[11] LightSeq Layer Normalization, https://github.com/bytedance/lightseq/blob/master/kernels/transformerKernels.cu.cc#L269
[12] cuBLAS, https://docs.nvidia.com/cuda/cublas/index.html
[13] GPT2,“Language Models are Unsupervised Multitask Learners”

字节跳动开源序列推理引擎lightseq(代码片段)

背景介绍2017年Google提出了Transformer[1]模型，之后在它基础上诞生了许多优秀的预训练语言模型和机器翻译模型，如BERT[2]、GPT系列[13]等，不断刷新着众多自然语言处理任务的能力水平。与此同时，这些模型的参数... 查看详情

字节跳动开源数据集成引擎bitsail的演进历程与能力解析

导读BitSail是字节跳动开源数据集成引擎，支持多种异构数据源间的数据同步，并提供离线、实时、全量、增量场景下全域数据集成解决方案，目前支撑了字节内部和火山引擎多个客户的数据集成需求。经过字节跳动各大业务线... 查看详情

一经开源就火了，字节跳动开源微服务中间件cloudwego(代码片段)

...整编综合自：https://github.com/cloudwego/kitex9月8日，字节跳动正式宣布开源CloudWeGo。这是一套以Go语言为核查看详情

去了字节跳动，回来聊一聊字节跳动的面试...(代码片段)

一、算法题一面：1.lc里最长上升子序列的变形题2.实现输入英文单词联想的功能二面：1.矩阵旋转，要求空间复杂度O(1)2.无序的数组的中位数。要求时间复杂度尽可能的小二、计算机网络tcp怎么保证数据包有序主机每... 查看详情

字节跳动开源最新gan压缩算法，算力消耗可减少至1/46(代码片段)

字节跳动近期开源了一项代号为OMGD的压缩技术。这是字节自研的GAN（生成对抗网络）压缩算法，在保证生成效果不变的前提下，算力消耗最低可以减少到原来的1/46，相比之前业界的最佳压缩效果提升一倍多。... 查看详情

揭秘字节跳动埋点数据实时动态处理引擎（附源码）(代码片段)

...场嘛，其中比较令大家感兴趣的就是最后一讲，字节一站式埋点平台的flink标准化清洗及拆流任务。其中大家感觉比较流啤的就是的就是字节做到了：不重启任务可以上下线新的拆流及清洗规则，所有的规则变更... 查看详情

字节模型被删？其实他们早有准备，lightseq上新支持transformer全流程训练加速，最高加速3倍！...

...训练时间成倍增加。针对这一痛点，字节跳动推出了LightSeq训练加速引擎，对Transformer训查看详情

字节模型被删？其实他们早有准备，lightseq上新支持transformer全流程训练加速，最高加速3倍！...

...训练时间成倍增加。针对这一痛点，字节跳动推出了LightSeq训练加速引擎&#x 查看详情

apachepulsar在火山引擎emr的集成与场景(代码片段)

更多技术交流、求职机会，欢迎关注字节跳动数据平台微信公众号，回复【1】进入官方交流群近年来，基于云原生架构的新一代消息队列和流处理引擎ApachePulsar在大数据领域发挥着愈发重要的作用，其应用场景和客户案例... 查看详情

揭秘字节跳动云原生sparkhistory服务uiservice

本文是字节跳动数据平台数据引擎SparkSQL团队针对SparkHistoryServer(SHS)的优化实践分享。*文|字节跳动数据平台—数据引擎—SparkSQL团队*在字节跳动内部，我们实现了一套全新的云原生SparkHistory服务——UIService，相比开源的SHS，UIServ... 查看详情

字节跳动开源bitsail，助力企业走好数字化“第一步”

10月26日，字节跳动宣布开源BitSail数据集成引擎。BitSail意为“数据航行”，可支持二十余种异构数据源间的数据同步，提供离线、实时、全量、增量场景下的全域数据集成解决方案，从而打通困扰企业数字化转型... 查看详情

字节跳动业务在flutter轻量级引擎上的实践与优化(代码片段)

本文介绍了字节业务在Flutter轻量级引擎上的实践历程，介绍了在此过程中遇到的各种各样的问题以及最终使用的解决方案。作者：字节跳动终端技术——候华勇一、背景Flutter在2.0版本之前混合工程开发对视图级别的开发... 查看详情

字节跳动2023秋招研发第五场笔试客户端方向(代码片段)

文章目录T1求和为0的最长连续序列思路代码T2族谱还原思路代码T3二次方程租思路代码T4AB实验室同学冲刺思路代码T1求和为0的最长连续序列通过：100%思路定义一个pre[i]表示从第一个元素到第iii个元素的影响和即前缀和思想... 查看详情

字节跳动2023秋招研发第五场笔试客户端方向(代码片段)

124.二叉树中的最大路径和-字节跳动高频题(代码片段)

一、题目描述路径被定义为一条从树中任意节点出发，沿父节点-子节点连接，达到任意节点的序列。同一个节点在一条路径序列中至多出现一次。该路径至少包含一个节点，且不一定经过根节点。路径和是路径中各... 查看详情

字节跳动2018校招算法方向（第一批）——2-最大值区间(代码片段)

时间限制：C/C++3秒，其他语言6秒空间限制：C/C++128M，其他语言256M给定一个数组序列,需要求选出一个区间,使得该区间是所有区间中经过如下计算的值最大的一个：区间中的最小数*区间所有数的和... 查看详情

字节前端终于开源！吹爆！(代码片段)

...，前端同学的福音大家好，我是鱼皮。最近，字节跳动的抖音前端技术团队开源了一款企业级应用设计系统SemiDesign 。这也是他们团队在GitHub上首次公开的项目，短短几天，就收获了3.6k个star。GitHub开源仓库老... 查看详情

字节跳动还是字符跳动(代码片段)

字节跳动还是字符跳动很多人对于java的io流的概念不太理解，什么是流呢？它其实就是指的是能够产生数据和接收数据的对象。那么有字节流的操作和字符流的操作，我们分别看一下这两种面向字节流的输入输出InputStream和OutputS... 查看详情