让算力不再成为瓶颈，小红书机器学习异构硬件推理优化之道
近些年，机器学习领域的视频、图像、文本和推广搜等应用，其模型计算量和参数量远远超过 CPU 摩尔定律的增长速度。在此背景下，GPU 的算力发展和大模型的发展不谋而合。很多公司都在结合 GPU 的算力发展，探索出适合自己的机器学习问题解决方案。例如，小红书在 2021 年开始进行推广搜模型的 GPU 化改造，以提升推理性能和效率。在迁移过程中，我们也面临一些困难，例如如何平滑迁移到异构硬件，如何结合小红书的业务场景和在线架构发展出自己的解决方案等等。在全球降本增效的趋势下，异构计算成为了一种很有前途的方向，可以通过将不同类型的处理器（如 CPU、GPU、FPGA 等）组合在一起来提高计算性能，从而实现更好的效率和更低的成本。
小红书推荐、广告、搜索等主要场景的模型服务，统一由中台推理架构承载。随着小红书业务的不断发展， 推广搜等场景的模型规模也在不断增大。以主推荐场景精排的主模型为例， 从 2020 年初开始，算法推出了全兴趣建模，用户历史行为记录长度均值扩大了约 100 倍。模型结构也从最初的 muti-task 经过多轮迭代，模型结构复杂度也不断提升 ，这些变化导致模型推理的浮点运算数增加了 30 倍，模型访存增加了约 5 倍。
模型特点：以小红书 2022 年底的推荐主模型为例，该模型具有充分的稀疏性， 部分结构由连续值特征和矩阵运算构成， 也存在大规模的稀疏参数比如 < 笔记 id * 用户城市 >，单个模型的 sparse 特征多达 1TB，但通过比较有效的模型结构优化，dense 部分控制在 10GB 以内，可放在显存中。用户每刷一次小红书，计算的总 FLOPs 达到了 40B， 超时的控制在 300ms 以内 ( 除去特征处理，带 lookup ) 。
推理框架：在 2020 年之前，小红书采用 TensorFlow Serving 框架作为在线服务框架，2020 年后，逐渐迭代成基于 TensorFlowCore 自研的 Lambda Service 服务。TensorFlow Serving 在进图之前进行一次内存拷贝 TensorProto -> CTensor，以确保模型推理的正确性和可靠性。然而，随着业务规模的扩大，内存拷贝操作会对模型性能产生影响。小红书自研框架通过优化免去一次不必要的拷贝，同时保留 Runtime、图调度能力、优化能力可插拔的特点，并为后期 TRT、BLADE、TVM 等不同优化框架的配合使用奠定了基础。现在看来，在合适的时候选择自研是一个明智的选择， 同时为了最大化减少数据传输带来的成本， 推理框架还承担了一部分特征抽取和转化的实现，这里小红书还在预估服务近侧部署自研的边缘存储， 解决了远端拉取数据的成本问题。
机型特性：小红书没有自建机房， 所有机器采购自云厂商，因此，选择不同机型的决策很大程度上取决于能够采购到什么型号的机器。而模型推理的计算并不是纯粹的 GPU 计算，合理找到硬件配比，除考虑 GPU\CPU 外，还涉及带宽、内存带宽、跨 numa 通信延迟等问题。
GPU 特性
GPU 特性：在这里，小红书和其它公司遇到的问题是一样的，GPU kernel 的执行可以分为以下几个阶段：数据传输、kernel 启动、kernel 计算和结果传输。其中，数据传输是将数据从主机内存传输到 GPU 内存；kernel 启动是将 kernel 代码从主机端传输到 GPU 端，并在 GPU 上启动 kernel；kernel 计算是实际执行 kernel 代码计算结果；结果传输是将计算结果从 GPU 内存传输回主机内存。如果大量时间都花费在数据传输和 kernel 启动上，而交付给 kernel 计算的活不重，实际计算时间很短，则会导致 GPU 的利用率无法提升，甚至出现空跑的情况。
预估服务框架
3.1 系统优化
3.1.1 物理机
在物理机优化方面，可以采用一些常规的优化思路， 主要目的是降低除 GPU 以外的其它系统开销成本， 降低虚拟化的中间商赚差价。一般来说，一套系统优化可以提升 1%-2% 的性能，从我们实践来看，需要结合云厂商实际能力来进行优化。
● 中断隔离：将 GPU 的中断单独分离出来，避免因为其他设备的中断而影响 GPU 的计算性能。
● 内核版本升级：提高系统的稳定性和安全性，提高 GPU 驱动程序的兼容性和性能。
● 指令透传：将 GPU 的指令直接透传到物理设备上，加速 GPU 的计算速度。
3.1.2 虚拟化和容器
在多卡情况下，将单个 pod 绑定到特定的 NUMA 节点上，从而提高 CPU 和 GPU 之间的数据传输速度。
● CPU NUMA Affinity，亲和性是指从 CPU 角度来看，哪些内存访问更快，延迟更低。如前所述，与该 CPU 直接相连的本地内存是更快的。因此，操作系统可以根据任务所在 CPU 来分配本地内存，以提高访问速度和性能，这就是基于 CPU NUMA Affinity 的考虑，尽量让任务运行在本地的 NUMA Node 里。在小红书场景里，CPU 上面的访存开销并不小。能够让 CPU 直连本地内存可以节约大量的 CPU 上 kernel 执行的耗时， 从而给 GPU 留足够的空间。
● 将 CPU 使用率控制在 70% 下，可以将延迟由 200ms -> 150ms。
3.1.3 镜像
编译优化。不同 CPU 对指令级的支持能力是有差异的， 不同云厂商采买的机型也有所不同。一个比较简单的思路是在不同的硬件场景下，编译镜像时带上不同的指令集合。在实现算子时，大量的算子本身已经带有如 AVX512 等指令。以阿里云的 Intel(R) Xeon(R) Platinum 8163 + 2 A10 的机型为例，我们根据该机型的特点和支持的指令集，编译优化调整合适的指令集， 整体相比于不进行指令优化的情况下，在该机型上的 CPU 吞吐量提高了 10% 。
3.2 计算优化
3.2.1 充分使用算力
● 计算优化，首先需要充分了解硬件性能，将其吃透。在小红书的场景中，如下图所示，我们遇到了两个核心问题：
1. CPU 上的访存较多， 内存 page fault 频率较高，导致 CPU 资源浪费，以及请求 latency 过高
2. 在线推理服务中，计算通常具有两个特点：单次请求的 batch size 小，单个服务的并发规模大。小 batch size 会导致 kernel 无法充分利用 GPU 的计算能力。GPU kernel 执行时间一般较短，无法充分掩盖 kernel launch 的开销，甚至 kernel launch 的时间比 kernel 执行时间还长。在 TensorFlow 中，单个 Cuda Stream launch kernel 成为瓶颈，导致推理场景下 GPU 利用率只有 50% 。此外，对于小模型场景（简单的 dense 网络），用 GPU 替换 CPU 完全不划算，限制了模型的复杂度。
● 为解决上述两个问题，我们采取了以下措施：
1. 针对内存 page fault 频率高的问题，我们使用 jemalloc 库来优化内存回收机制，并开启了操作系统的透明大页功能。此外，针对 lambda 特殊的内存访问特点，我们设计专门的数据结构，并优化内存分配策略，尽可能地避免内存碎片。与此同时，我们直接绕开了 tf_serving 的接口，直接调用 TensorFlow，减少了一次数据的序列化与反序列化。这些优化在首页精排和内流精排场景下，提升了 10+% 的吞吐，在广告大多数场景下，降低了 50% 的 latency。
兼容 tensorflow::Tensor 格式，在将特征传递给 tensorflow::SessionRun 之前是零拷贝
2. 针对 TensorFlow 单 Cuda Stream 的问题，我们支持了 Multi Streams , Multi Contexts 的功能，避免了互斥锁导致的性能瓶颈，成功将 GPU 利用率提升到 90+% 。同时，我们利用 Nvidia 提供的 Cuda MPS 功能，实现了 GPU 的空分复用 (同一时间支持多个 kernel 执行)，使得 GPU 的利用率进一步提升。基于此，Search 的排序模型成功在 GPU 上实现。此外，我们也在其他业务线上成功落地，包括首页初排、广告等等。下表是在搜索排场景下的一个优化情况。
3. Op/Kernel fusion 技术：通过手写或者图编译优化工具生成性能更高的 Tensorflow 算子，充分利用 CPU 的 Cache 以及 GPU 的 Shared Memory，提升系统的吞吐。
在内流场景下， 算子进行融合，可以看到单次调用 12ms -> 5ms
3.2.2 避免算力浪费
1. 系统链路上存在优化空间
a. 初排前置计算：在处理用户侧相关计算时，初排需要计算大量笔记，例如以外流为例，需要计算约 5000 篇笔记，lambda 对其有切片处理。为避免重复计算，将初排的用户侧计算前置到和召回阶段并行，从而使得用户向量的计算从多次重复变成了只需要 1 次，在粗排场景下优化了 40% 机器。
2. 图内训练到推理过程中：
a. 计算前置：通过 graph freeze 可以将一部分计算提前处理。在推理时，不需要重复计算。
b. 产出模型 freeze 优化：模型产出时把所有的参数和图本身一起生成冻结图（ frozen graph ）并进行预处理计算，可以将很多预计算的 Variable 算子转换成 Const 算子（ GPU 使用率下降 12% ）
c. 推理场景下的合并计算：每个 batch 只包含一个 user ， 即用户侧存在大量重复计算， 具备合并的可能性
d. CPU/GPU 算子拆分：将 lookup 之后的全部算子移至 GPU ，避免了 CPU 和 GPU 之间的数据拷贝
e. GPU 到 CPU 数据拷贝：将数据打包一次拷贝
f. BilinearNet 算子 GPU cuda 实现：通过 GPU 加速计算，提升性能
g. 部分算子 GPU 化：省去 CPU -> GPU 拷贝
h. BatchNorm & MLP 合并：通过实现新的 MLP 层，根据一个目标减少进 GPU 的次数 （ N -> 1)， 增大一次计算的计算量（重复利用 GPU 小核心的并发能力)
3.2.3 全天动态算力
● 动态计算降级提升全天资源使用效率，秒级别的对 lambda 负载进行自动负反馈调整，做到对单区压测之前不需要人工做降级准备。
● 在外流精排、外流初排、内流精排、内流初排、搜索等主要业务场景均已经上线。
● 在多个业务线解决了容量问题，有效缓解了业务增长导致的资源线性上升，同时大幅提升了系统的鲁棒性。在功能上线后的业务线中，均没有出现因为瞬间成功率大幅下降导致的 P3 及以上事故。
● 大幅提升全天资源使用效率，以内流精排为例（如下图所示），五一假期三天的 10:00-24:00 的 CPU 使用核数均保持 50 核的一条平线（抖动对应发版）
3.2.4 换更好的硬件
● A10 GPU 的性能是 T4 GPU 性能的 1.5 倍，同时 A10 机型配备的 CPU ( icelake, 10nm ) 比 T4 机型 ( skylake, 14nm ) 更新一代，价格仅为 T4 机型的 1.2 倍。未来我们还会考虑在线使用 A30 等机型。
3.3 图优化
3.3.1 DL 栈的自动编译优
● BladeDISC 是阿里最新开源的基于 MLIR 的动态 shape 深度学习编译器， 小红书的自动图优化部分来自于这套框架（ Blade 推理加速库是 Apache 2.0 开源，可以跨任何云使用，无知识产权风险）。该框架提供了 TF 图编译优化（包含 Dynamic Shape Compiler ，稀疏子图优化），同时能叠加我们本身做的算子定制化优化，可以较好的适配我们的业务场景。在压测单机 inference 中，QPS 能提升 20% 。
● 这套框架关键技术
（1） MLIR 基础架构
MLIR，即多层次中间表示语言（Multi-Level Intermediate Representation），是由 Google 发起的开源项目。其目的是提供一个灵活、可扩展的多层 IR 基础设施和编译器实用工具库，为编译器和语言工具的开发者提供一个统一的框架。
MLIR 的设计受到 LLVM 的影响，但与 LLVM 不同的是，MLIR 主要关注于中间表示（ IR ）的设计和扩展。MLIR 提供了一个多层次的 IR 设计，可以支持从高层语言到底层硬件的编译过程，并提供了丰富的基础设施支持和模块化设计架构，使得开发者可以很方便地扩展 MLIR 的功能。此外，MLIR 还具有较强的胶水能力，可以与不同的编程语言和工具进行集成。MLIR 是一个强大的编译器基础设施和工具库，为编译器和语言工具的开发者提供了一种统一的、灵活的中间表示语言，可以方便地进行编译优化和代码生成。
（2） 动态 shape 编译
静态 shape 的限制意味着在编写深度学习模型时需要提前确定每个输入和输出的形状，并且不能在运行时改变它们。这限制了深度学习模型的灵活性和可扩展性，因此需要一种支持动态 shape 的深度学习编译器。
3.3.2 精度调整
● 量化的实现方式之一是使用 FP16
FP16 计算优化：在 MLP 层时用 FP16 替换 FP32 计算，能够较大地减少 GPU 使用率(相对下降 13% )
在调整 FP16 的过程中，选择白盒方式进行精度优化意味着可以更加精细地控制哪些层使用低精度计算，并能够根据经验进行不断调整和优化。这种方式需要对模型结构有较为深入的了解和分析，可以根据模型的特性和计算要求进行有针对性的调整，以达到更高的性价比。
相比之下，黑盒方式则相对简单，不需要了解模型的内部结构，只需要设置一定的容忍阈值即可完成精度优化。这种方式的优点是操作简单，对模型同学的要求也相对较低，但是可能会牺牲一定的性能和精度。
因此，选择白盒还是黑盒方式进行精度优化需要根据具体情况而定。如果需要追求更高的性能和精度，同时拥有足够的经验和技术能力，那么白盒方式可能更加适合。如果操作简单、快速迭代更加重要，那么黑盒方式可能更加实用。
从 2021 年开始到 2022 年底， 经过本项目优化，小红书推理计算算力增加 30 倍，关键用户指标提升 10%+ ，同时累积节约集群资源 50%+ 。在我们看来，小红书在 AI 技术方面的发展路径应该是以业务需求为导向，平衡技术和商业的发展：实现技术创新的同时，也要考虑成本、效益和可持续性。以下是一些优化过程中的思考：
优化算法和提高系统性能。这是小红书机器学习团队的核心任务。优化算法和提高系统性可以更好地支持业务需求，提高用户体验。然而，在资源有限的情况下，团队需要明确优化的重点，避免过度优化。
构建基础设施和提高数据处理能力。基础设施对于支持 AI 应用是非常关键的。小红书可以考虑进一步投入基础设施的建设，包括计算和存储能力、数据中心和网络架构等。此外，提高数据处理能力也是非常重要的，可以更好地支持机器学习和数据科学应用。
提高团队人才密度和组织架构。一个优秀的机器学习团队需要拥有具有不同技能和背景的人才，包括数据科学家、算法工程师、软件工程师等；优化组织架构也有助于提高团队效率和创新能力。
合作共赢和开放创新。小红书持续与其他公司、学术机构和开源社区合作，共同推进 AI 技术的发展，这有助于小红书获取更多资源和知识，成为更加开放和创新的组织。
该方案让小红书机器学习架构水平达到了业界一流水平。未来，我们将不断推进引擎升级和降本增效， 引入新技术提高小红书机器学习的生产力， 将更加结合小红书的实际业务场景， 从单模块的优化升级为全系统优化，并进一步引入业务侧流量的个性化差异特征， 将降本增效做到极致。期待有志之士，一同加入我们！
张楚岚（杜泽宇）：商业技术部
毕业于华东师范大学， 商业化引擎团队负责人， 主要负责商业化在线服务搭建。
陆光（彭鹏）：智能分发部
毕业于上海交通大学 ，机器学习引擎工程师，主要负责 Lambda GPU 优化。
伊恩（陈建新）：智能分发部
毕业于北京邮电大学，机器学习引擎工程师， 主要负责 Lambda 参数服务器和 GPU 优化。
赤羽（刘兆宇）：智能分发部
毕业于清华大学，机器学习引擎工程师， 主要负责特征引擎方向的相关研究和探索。
特别感谢 ：智能分发部 所有同学
2亿人生活方式分享背后的多模态学习 2022-04-11 加入
小红书技术团队官方账号，小红书技术创新与问题解读的分享平台，与你共前进。
评论