NLPCC：预训练在小米的推理优化落地

导读： 本文主要分享 小米 AI 实验室 NLP 团队 在 NLPCC 轻量级语言模型比赛 上的经验，以及我们在预训练模型推理优化上所作的工作和达到的实际落地后的效果。此次分享的目的是帮助大家快速进入比赛，以及了解工业界中使用 BERT 之类的大型预训练模型时进行推理的相关优化手段。

01 背景介绍

首先和大家说一下比赛的背景和预训练模型中存在的问题。

1. NLP 中的预训练

随着 BERT 的推出和发展，预训练模型成为目前工业界和比赛使用最广泛的模型。目前在各大 NLP 任务 SOTA 榜单上排前几名都是大型的预训练模型，比如原生 BERT 或者它的一些变体。

预训练模型的应用分为两个阶段，先进行预训练阶段 ( pre-training )，然后进行微调阶段 ( fine-tuning )。预训练阶段利用大规模的无监督预料，通常大于 100g 的量级，加上特定的无监督任务进行训练。用来预训练的任务有，经典的 NSP ( next sentence predict )、MLM ( masked language model )、以及一些变体，如阿里发布的 Structural language model。另外在预训练阶段时，embedding 层的使用方式也有很多，比如 NEZHA 以及 XLNET 使用的相对位置编码。Pretrain 阶段一般比较消耗计算资源，一般都是在大量的 GPU 集群上进行运算，并且计算很长的时间才能得到一个比较好的结果。

相对于复杂预训练任务，下游的微调任务就比较简单了。在预训练好的模型的基础上，只需要添加较少的网络参数，并在在下游有监督数据上进行少量的训练，就可以得到很不错的成果，这也是预训练技术能够被大量使用的基础条件。

2. 效率问题

然而就像天下没有免费的午餐一样，预训练技术在带来好处的同时也存在一些问题，这其中最为凸显的恐怕是效率问题了。

微软在发布 TuringNLG 模型的同时发布了这种图，从上图我们可以看到在 BERT 发布到现在，自然语言模型参数量的发展是指数型增大的，从 BERT 的 1 亿参数量、到 GPT2 的 10 亿、然后 TuringNLG 的 100 亿。模型变大带来的效率问题可以从训练和推理两个方面来体现。

在训练阶段，一个好的模型需要大量的语料支持，通常为了得到一个不错的模型，需要使用 100g 以上的训练数据。同时为了支持如此规模的数据集，使模型能拟合出不错的效果，就需要使用大量昂贵的 GPU 资源。而在推理阶段，大模型的推理时间比较长，那如何在高并发业务线上使用大模型，并且保证推理服务的性能 ( 比如服务延迟 99 分位数 ) 成为一个难题。由此可见，在预训练模型应用时必须考虑性价比，在效果提升和成本预算之间做一个合适的权衡。

3. 解决方向

针对这些效率上的问题，衍生出两个比较清晰的解决方向和思路。

第一种方法是小模型或蒸馏，小模型是指参数特别少的模型，比如在 NLPCC 小模型比赛中聚焦的小于 12M 参数的模型。蒸馏是将训练好的 BERT 类大模型的能力迁移到较小的模型中，目前业界聚焦的任务是将 12 层的 BERT-base 模型蒸馏到 6 层的 BERT 上，并保证较少的性能损失。这种方法的主要优势是资源占用率比较少，适用于线上的高并发业务；并且通过一些方法能够实现与大模型类似的效果。

第二种方法是对大模型的优化，具体来说是指优化大模型在单张 GPU 上所能承受的最大流量，即吞吐量。这种方法的优势是效果较好，比较适合低并发业务。

02 NLPCC 预训练小模型比赛经验分享

1. 比赛介绍

比赛要求的小模型必须是小于 12M 参数的预训练模型，即 1/9 的 BERT-base 大小，最后在四个下游任务中测试模型效果并进行打分。评测涉及任务有：指代消歧 ( CLUEWSC2020 )、分类任务 ( CSL 论文关键词识别 )、命名实体识别 ( CLUENER2020 ) 以及阅读理解 ( CMRC 2018 )。其中，指代消歧任务用的是小数据集 ( 1.2k 左右 )，其余任务所用数据集大小为万级别。在小数据集上微调模型，尤其是小预训练模型，是比较影响模型效果的，后面我们会介绍这方面所做工作，希望能够对读者有所帮助。

用一句话来总结这个比赛，小模型的效果真的很差吗？我们直接来看比赛的结果。

其中，第一条我们提交的 BERT-base 的参考。可以看到相较于 BERT 模型，小模型在分类、NER 以及阅读理解任务上都可以将差距缩小到 1.5%以内。而在 wsc 任务上由于数据集太小，导致微调时效果的抖动非常大效果会略差一点，而实际中很少有这么小的数据集。

2. 经验分享

接下来主要介绍一下我们的解决方案。在直接预训练小模型和蒸馏的方式选择上，我们采用了前者。

我们在比赛中所做的工作主要分为了 3 个方面，首先是数据。我们一共收集了 160g 的原始中文语料，并剔除了一些口语、表情、微博段子、网络流行用语比较多的语料，最终筛选出 35g+相对通用的语料。为了快速得到实验结果，我们在输入长度上采用了 256 的大小，同时也正在尝试 512 的长度，以便在阅读理解等需要更长的句子输入的任务上得到更好一点的效果。

在模型方面，我们选用了 6 层的高瘦模型，并通过以下公式选择对应的 Hidden size 和 Vocab size 来保证模型参数总量小于 12M。

大概介绍一下训练细节，我们使用了 8 张 V100 卡；采用了混合精度的方法；优化器为 lamb 优化器；batch size 大概为 14400，并每过十几步累积一次梯度；预训练任务选用了 wwm-MLM。没有选用 NSP 任务的原因是，我们在测试 MLM + NSP 任务训练时出现了收敛较慢的现象，原因仍有待分析。最终训练结果是 MLM loss 达到了 1.5 ~ 1.8 之间，同时我们发现当 loss 到 1.8 左右时，分类任务和 NER 任务对应的表现已经很好了；继续训练当 loss 来到 1.5 附近时，阅读理解任务的效果提升较大，而分类任务和 NER 任务的表现开始有一些抖动。

遇到的问题：

同时我们分享两个在比赛过程中遇到的问题。

第一个是模型训练时的梯度爆炸，具体是每当 loss 收敛到 1.5 附近时就会出现梯度爆炸的问题。再排除了数据中的问题后，最终我们找到的解决方法是将初始的学习率调小，之后再没有遇到同样的问题。

第二个问题是，我们再尝试微软发布的 Bert of theseus 蒸馏方式时，发现该方法在本次小模型比赛中取得的效果不好，原因是使用该方法时 Hidden Size 不能改变，因此必须压缩层数来达到参数量的限制。同时，我们尝试用小的 Hidden Size 预训练一个模型，再去蒸馏一个 12M 的模型。可能是受限于 Teacher 模型的效果限制，得到的结果是不好的。

3. 后续尝试

同时，在比赛快结束的时候我们总结性的提出一些后续可以继续尝试的思路。

注：NLPCC 比赛结束后，我们又尝试了一个数据增强和简单的蒸馏方法，目前在 WSC 和 CLUENER 两个任务上已经超过了 bert-base 的效果，并且排在了小模型排行榜第一名，详见如下：

以上就是这次 NLPCC 小模型比赛的经验分享，下面将分享一下大模型推理服务的优化以及在小米的落地情况。

03 大模型推理效率及在小米的优化落地

1. 相关知识和现状

为了方便大家阅读，我们先介绍两个工程上评价服务的指标。

为什么要介绍这两个指标，是因为线上服务要在短时间内接受大量的请求，为了满足使用者的体验以及维护公司利益，必须要保证较高的可用性。

我们再来看一下在用大模型推理时，用一张 T4 卡能达到的最大 QPS。如果模型用的是 BERT base，句子计算长度设定为 16，那么最大的 QPS 大概是 100，这时 P99 已经比较高了，如果 QPS 设定为 150 服务将变得不可用；同时，如果模型是 BERT-large，在同样的句子长度下，模型根本无法运行。在这种情况下假设有一个日活 1000w+的服务要上线，QPS 对应要求为 2000，并使用 BERT-base 作计算推理。那么我们计算了一下，一共需要 20 张 T4 卡来满足需求，并且 P99 延迟会大于 90ms，这是一个非常大的资源浪费。

2. 优化效果

开门见山的说，在优化后同样一张 T4 卡上可以达到什么样的效果呢？大家可以看下表中 QPS 和对应的 P99 表现，当 QPS 最大达到 3000 时，P99 的延迟仍然保持在 40ms 以下。

如果大家为了追求更好的效果使用 BERT-large 模型，使用相同的推理优化方法能够达到的数据如下表，极限 QPS 大概是 800 左右。

总结一下优化效果，优化前如果 BERT-base 模型要放在线上推理，需要使用 20 张 T4 来满足需求。优化后可以将数量缩减到 1 张，一共节约了 19 倍的成本。

3. 实现方案

下面来介绍一下我们的优化方案。

简单来说分为三个：

• 第一个是Tensorflow Serving提供的Batching功能，该方法能够提升n倍的服务性能；

• 第二个是FP16的量化，可以提升2倍服务性能；

• 第三个是NVIDIA开源的Faster Transformer库。

Tensorflow Serving Batching 所做的是将多个请求合并在一起进行计算，这样可以利用 GPU 的并行性能。之所以能这样做是因为 GPU 在计算一定量的 batch 所消耗的时间和计算单挑所用时间相差不多。具体的原理是服务端在规定好的一段时间里持续接收请求并放入队列等待，如果请求数达到规定的数量或者队列等待时间超过规定时间，再进行推理运算。另外一个经验是在使用 TF-Serving Batching 时，服务端做 Padding 会用随机数填充，所以建议大家在 Client 端做 Padding。

在用 FP16 优化中，我们使用了 TensorCore 来做加速，将模型图中 FP32 精度的节点全部转换为 FP16，从而降低模型推理需要计算的运算量 ( 每秒浮点运算次数 TFLOPs )。在 BERT-base 上实验，这种方法很稳定，损失基本上控制在万分位，是一个非常值得使用的方法。

第三个方法是用 Fast Transformer，大家感兴趣可以去 github 了解源码。我们在这基础上所做的工作是，将 Fast Transformer 集成到 Tensorflow Serving 上，并且支持对可变长度输入的处理等等。大家可以访问英伟达的 github 主页了解更多。

4. 线上效果

前面在评测这些优化方法时用到的数据都是压测数据，那么在真实线上业务的表现怎么样呢，我们用小爱同学的闲聊业务来举例说明。

这个任务是一个问答排序问题，比如用户问"你喜欢维尼小熊嘛"，这时对应的答案可能有好多种，按照先检索后排序的方法，需要对输入和每个回答计算一个相似度。我们使用的方案是 Pointwise 的 BERT 二分类，使用的模型为 6 层的 BERT，输入长度为 32，输入为 Query+Reply，输出为 0-1 的相关性。下面是服务的优化效果，我们一共分为了三个版本，在使用了这三种优化方法后成功的将 13 张 GPU 压缩到了 3 张 GPU，并且 P99 从 200ms+降到了 35ms，效果还是比较明显的。

04 总结

我们的分享就结束了，一共讲了三个部分，先分享了我们在 NLPCC 小模型比赛上的经验，然后介绍了预训练推理优化在小米的相关落地，最后感谢赛事主办方 CLUE&NLPCC 提供这样一次赛事。

今天的分享就到这里，谢谢大家。

作者介绍：

赵群，小米高级软件工程师

2014 年硕士毕业于中国人民大学，同年加入微软 Bing 广告组。2017 年加入小米 AI 实验室 NLP 应用组，参与了闲聊机器人在小米小爱同学的搭建与落地，现聚焦在预训练通用优化方向。

本文来自 DataFunTalk

原文链接：

https://mp.weixin.qq.com/s?__biz=MzU1NTMyOTI4Mw==&mid=2247502777&idx=2&sn=ebc2ccc9997a45732b4244cfbb35f63f&chksm=fbd779d5cca0f0c3ecca9fa6d2c736e8f1c497116c2157e0da23f0db19fb46df2f84b5096fbc&scene=27#wechat_redirect