文丨阿里云神龙计算平台AI加速团队 & UC搜索架构部推理引擎团队
导语：作为国产行列里占有率排名第一的移动浏览器，UC浏览器自身承载着数以亿计的用户量，当前UC浏览器每天的服务请求对服务器的算力及带宽要求极高，因此也带来了巨额的运营成本。因为业务是动态变化的，UC对计算资源也有动态扩缩容的需求。阿里云GPU云服务器是提供GPU算力的弹性计算服务，具有超强的计算能力，服务于深度学习、科学计算、图形可视化、视频处理多种应用场景，能为客户提供软件与硬件结合的完整服务体系，助力客户在实际业务中实现资源的灵活分配、弹性扩展、算力的提升以及成本的控制。而基于阿里云GPU云服务器的神龙AI加速引擎（AIACC）是为了极致性能而生，旨在打造世界级无与伦比的AI性能体验，同时为了客户降本增效，实现共赢。据悉，刚公布的最新世界MLPerfTM推理榜单中，基于阿里云GPU云服务器的AIACC首次突破封闭式场景世界第一。
本篇文章将带大家了解阿里云AIACC如何基于阿里云GPU云服务器助力UC浏览器的搜索业务平衡计算性能与运营成本之间的矛盾，使其大幅实现降本增效，成功在阿里云的GPU云服务器落地。
背 景
1. 业务背景
UC搜索承载着UC主要业务入口，场景包括：大搜、各种垂搜业务、夸克app等。搜索流程一般经过几个阶段：召回 –> 粗排 -> 精排(L2->L4) -> 混排等。架构如下：
在业务中，L3/L4排序部分都使用了QTC核心模型。随着业务流量的增长，目前精排打分阶段面临巨大挑战，延迟和吞吐都要兼得。
2. QTC模型
下图是用TF-summary对QTC模型做可视化，
QTC模型属于排序核心模型，模型结构分为3个BERT+ 多层Conv + 多层MLP等，其中Bert输入seq length最大长度是512。模型总共有大约4500个算子，计算量巨大。
3. 原始性能
最初采用了NV提供的Faster-Transformer这一套软件来优化QTC模型的推理，但由于Faster-Transformer只能加速BERT网络，导致推理时需要做多个子图割裂运行，加上其余网络部分也未做优化，从而整体性能并不好。针对A100机型，做了基于这套方案的性能评估，延迟控制在30ms下，QTC模型只能跑到350QPS。
4. 挑战与目标
随着模型效果带来的业务收益，促使业务流量不断增加，在算法预算前提下，扩量面临比较大的挑战。一方面资源不够，一方面资源也没有完全利用起来导致的吞吐不够，同时观察分析在GPU使用率达到90%以上时延迟会出现飙升的情况，这种情况对于稳定性带来了不小的影响。如何在原来预算的基础上提高吞吐、降低延迟、提升性价比最终规模化上线为业务赋能成为最大的挑战。
结合搜索业务的特点，针对QTC模型在GPU上推理提出了几个方面的需求：
1. 为了保证算法的效果，推理时模型输出的误差控制在千分位；
2. 为了保证终端用户在搜索时的体验，精排部分的推理耗时在30ms以内；
3. 在有限的预算下满足业务的需求，即QPS需要远高于原有解决方案；
综合业务方各个方面的需求，UC制定如下的目标：
在A100机型上，保证精度的前提下，QTC模型耗时在30ms以内，QPS达到1120以上。
优 化
1. 优化路线
TensorRT是Nvidia开发的针对模型推理的高性能软件库，但QTC模型复杂度高，无法整体直接用TensorRT load。此外，直接用TensorRT load时部分层会被拆成细粒度的op，性能较差。对QTC模型做深度的分析之后，结论是用TensorRT以及其plugin支持的op可以覆盖QTC模型的op，但是现有的parser无法按照需要的粒度去解析QTC模型并映射到TensorRT上。
针对上述问题，AIACC从模型定义层，直接解析QTC模型并映射到TensorRT上，避免转PB后大的operator被拆分成细粒度的operator，以达到最佳性能。此外，这么做也便于以不同的粒度进行模型精度的分析，以不同的粒度实现op的融合与替换，方面集成AIACC的高性能Kernel进来。
2. Enable TensorRT & 精度对齐
2.1 映射正确性验证
解析QTC模型并映射到TensorRT上的过程中，首先要解决的问题是映射的正确性验证。对正确性验证，首先要解决基准是什么这个问题。我们拿到的信息为一份模型结构的代码与一个训练中的checkpoint。基于业务方给的有限的信息，完成了从训练框架侧load checkpoint的逻辑，并根据需求，留下了可以获取任意中间结果的接口。在验证映射正确性的过程中，我们以训练框架侧运行推理的结果为对比的baseline，和我们构建的TensorRT侧的Engine的运行结果做对比，来判断是否正确的把QTC模型映射到TensorRT上。
下面以一个self-attention层为例，展开怎么做QTC模型到TensorRT的映射并保证运算结果的一致。
在TensorRT侧，为了校验映射的正确性，需要有推理的脚本来验证build好的engine的正确性。同在训练框架侧构建正确性的baseline类似，我们基于TensorRT开发了一套load build好的engine并做推理的脚本，来获取engine的推理结果，用于和训练框架侧的baseline做对比。
TensorRT为了追求更高的性能，对self-attention层的运算做了等价变化，把部分转换提前到模型build阶段，进而减少模型推理需要的operation。其他的一些细节，包括mask处理、数据类型、hidden size等，按照类似逻辑，一一对齐后，我们即可在TensorRT侧构建出运行结果与训练框架侧运行结果一致的engine。
对比TensorRT直接load checkpoint的模式，在我们的映射中，把self-attention映射为1个op，避免了拆成多个细碎op而导致的性能问题。这种模式下self-attention的信息在一个节点上，也方便我们后续做kernel的优化，例如融合self-attention内部的GEMM等多个操作，以达到更好的新能。
2.2 数值问题解决
数值问题是AI任务中很难定位解决的一类问题，因为没有编译器或者其他的报错提示来帮助我们定位问题。在把QTC模型映射到TensorRT的时候，遇到了数值问题，现象为多个相同的基础模块叠加导致中间的feature map中有NAN的异常值，进而导致最终的结果误差远远超出业务团队要求的千分位。
第一时间把这个问题定为数值问题，是因为在QTC模型到TensorRT映射过程中，每一个子模块我们都做了单元测试来验证映射的正确性，不存在映射错误导致的异常。此外，为了推理的性能，我们开启了FP16，部分中间的运算是用FP16进行，FP16的表达能力相比FP32差很多，这是引入数值问题的一个风险点。多个相同的基础模块叠加导致最终的输出中有NAN的异常值，大概率是因为计算中出现了极大or极小的数值，导致某些操作，例如做exp/除法的时候，出现的异常的行为。
分析清楚问题后，解决问题的方法也很简单，完善TensorRT中这部分exp计算逻辑，避免exp计算中出现大的负值即可。
2.3 其他的N个问题的概述
在把QTC模型映射到TensorRT并对齐精度的过程中，我们还遇到了其他一系列大大小小的其他问题，这里不一一展开，只在简单列举一下，包括用shuffle替换reduce来提高实现reshape的运算的效率，用identity层来实现data type cast的逻辑，conv层的参数转换，pooling的padding处理，reshape后接FC层参数排布等问题。
3. 关键Kernel优化
seq_length=35/512时，TensorRT未对Multi-Head Attention做针对性优化。seq_length为512的Bert耗时占比较大，如果对这部分做针对性优化，会获取较大幅度的性能提升。对这部分的优化我们分两步进行，第一步优化了seq_length=512的Multi-Head Attention中的softmax的实现，第二步则是针对seq_length为35/512的case做了类似TensorRT针对seq_length=64的情形下做的融合。
下图是seq_length为512的情况下，未优化前的一个Transformer层调用的Kernel，其中绿框中的kernel为第一步优化的Kernel，红框中的部分则是第二步优化中融合为一个。
4. Enable CUDA-Graph
CUDA Graph是一个把所有kernel算子结合（capture）成Graph，然后整体launch这个Graph，减少频繁的kernel launch来带的开销以及kernel中间的gap。下图是普通的kernel执行和Graph执行的区别。可以看出，在kernel执行时因为需要CPU和GPU切换，造成小算子间会有比较大的GPU idle时间（gap引起），同时如果小算子执行的时间比较短，那么launch的时间占比就成了大头，造成GPU利用率低下。
在CUDA11版本和Ampere架构下，CUDA Graph本身做了很大的改进，在CUDA内部做了并行化处理。
从机制图可以看到，CUDA内部在执行Graph时，查找依赖关系，如果发现无直接依赖且目前资源满足算子执行（比如：stream processor/share memory/register等）则并行执行，提高GPU利用率。同时我们也对CUDA Graph下做了一些优化（比如增加memory cache机制、graph min update、multi-stream处理等）更好地提升了性能。
Enable CUDA Graph有两种方式，其中流捕获提供了一种从现有的基于流的 API 创建图的机制。在QTC模型中，通过流捕获的方式来加速基于TensorRT构建的图。Enable CUDA Graph后，在A100上延迟有3%下降，吞吐提高了4%，在GPU使用率高时latency也更稳定。
业 务 结 果
下图中，Faster-Transformer是指UC原有的基于Faster-Transformer的在GPU上的解决方案；AIACC-Optimization是我们优化后的性能数据：
在UC的生产环境中，A100机型上QPS达到了1330，为最初Faster-Transformer版本的3.8倍，超出了预设的目标。
目前，采用AIACC这一系列优化的QTC模型已经在部分搜索业务上线使用，后续即将大规模线上使用。按最低的预期使用量来计算，每年能节省达数千万的成本。
在这次合作中，UC搜索精排应用在阿里云GPU云服务器上运行的性价比得到了显著提升，可以把业务批量部署在生态成熟的阿里云GPU云服务器上。云平台弹性扩缩容的优点让UC可以根据业务需求，按需从阿里云上获取GPU云服务器资源。此外，推理用的机型和训练用的机型统一，让生产链路上从算法到优化再到部署的工程师合作更方便，也方便后续在线上实现训练与推理任务的混合部署，最大化GPU资源的利用。
阿里云AIACC针对阿里云GPU云服务器硬件做了深入的优化，达成了在不额外引入新硬件平台的情况下，用阿里云GPU计算实例来满足UC极致性价比需求的目标。此外，用UC真实应用打磨的优化性能工作都沉淀到AIACC中，后续可以规模化服务更多的阿里云GPU云服务器的客户。
点击这里，了解阿里云GPU云服务器。
