文丨阿里云神龙计算平台 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 云服务器。