GPU 加速 Pinterest 推荐模型，参数量增加 100 倍，用户活跃度提高 16%

2022 年 8 月 10 日
重庆
本文字数：3486 字
阅读完需：约 11 分钟

作为图片届的“Twitter”，Pinterest 首页展示给用户的图片也离不开背后的推荐模型。近期，其工程团队通过将机器学习服务从 CPU 转移到 GPU 上，使得 Pinterest 可以上线比之前大 100 倍的推荐模型。上线大模型给模型质量带来了阶跃式的提升，最终将 Pinterest 首页 feed 流的用户活跃度提高了 16%。

在本文中，他们分享了如何只经过微小的投入和延迟的成本实现这一优化目标，这包括优化特定算子，整合内存传输，通过 CUDA Graph 技术在设备上执行静态图，以及分布式系统的重新设置。他们的实践证明，要想提高模型效果，模型参数就得变大，而使用 GPU 服务方案的经济效益远超 CPU 服务方案。

来源｜Pinterest Engineering

翻译｜郑泽康

Pinterest 的使命是给每个人带来创造他们所热爱生活的灵感。为了实现这一使命，我们所有产品中都包含的一个关键组件就是各式各样的推荐模型，它们负责在合适的时间给合适的人展示合适的内容。我们的推荐模型是通过高级算法进行训练的一系列机器学习模型，用于理解用户在 Pinterest 上花费时间时的行为，这些推荐模型是通过定制的机器学习模型服务器（Scorpion Model Server, 即 SMS）来运行的。

SMS 上需要面对十分艰巨的技术挑战，基于 3000 多亿个 Pin 的语料库，它必须在毫秒内为 4 亿多个用户提供相关推荐。之前，SMS 在 CPU 上进行模型推理，其内核经过了多年的优化以满足我们对延迟和基础设施成本的严格要求，但即使最新一代的 CPU 也几乎达到了 SMS 服务的极限。我们必须非常审慎地确保，每次因模型变化而带来的延迟和基础设施成本的增加都是合情合理的。

机器学习领域里模型参数和计算量增加的趋势让问题变得更加严峻。在推荐系统中，具有 1000 亿参数量的模型已经很常见，并在业内常被提及。

在 Pinterest，我们采用了稍微不同的方式，通过使用诸如 Transformer 的现代模型架构来扩大模型。在更大模型下，我们立即观测到模型准确率发生的质变——其大幅提升了 Pinner（译注：Pinterest 用户）的参与度。但是，在 CPU 服务器上运行这些现代模型架构几乎让成本和延迟都提升了 40 倍，这一代价难以承受。因此，我们转而寻求使用 GPU 来加速模型推理，从而可以使用合理的成本来运行这些模型。

1、优化

当我们尝试那些开箱即用的 GPU 服务时，很快意识到在经济高效地利用 GPU 运行推荐模型服务之前需要对其优化。我们首先使用分析工具来分析在模型推理过程中发生了什么，在仔细观察分析结果时，我们注意到时间线上有大量的小 CUDA Kernel 在执行。

这是推荐系统模型的预期行为，其中数百个特征在模型后期阶段进行特征拼接之前需要单独处理。然而，对于大量的小算子，启动 CUDA Kernel 带来的开销比计算开销还要昂贵。与训练时的 batchsize 相比，模型服务时的 batchsize 相对更小，进而加剧了这一问题。

上图分别是模型优化前和优化后的分析结果。CUDA Kernel 时间线（用红框突出显示）表明，Kernel 的启动开销（蓝色块之间的间隙）显著减少，因此 GPU 得到了更好得利用，它将大部分时间花费在 Kernel 执行中。

2、减少小 op 的数量

我们采取的第一个方法是找出减少小 op 数量的机会。我们寻找常被使用的模型组件，并尽可能对其优化。其中一个例子是 Embedding 的查表模块，它包含两个计算步骤：原始 ID 到索引的查找，索引到 Embedding 的查找。由于我们有大量特征，这两步操作会被重复数百次。通过使用 cuCollections (https://github.com/NVIDIA/cuCollections) 以支持 GPU 上原始 ID 的哈希表，并实现了自定义的 Embedding 查找模块以合并多个查找操作，我们显著地减少了 op 的数量。在执行一些优化后，马上看到了更优的性能表现。

3、合并内存拷贝

同样，当我们在主机和 GPU 显存之间搬运 Tensor 时，也存在整合数据传输的机会。通用推荐模型里的一个候选样例通常会使用数百个特征作为输入，对于每一次推理，各个特征作为一个单独的 tensor 被拷贝到 GPU 显存中。虽然在主机和 GPU 显存之间搬运数据非常快，但是为每个请求调度数百个 cudaMemcpy()函数调用的开销很快成为瓶颈。

从主机单独拷贝 Tensor 到 GPU vs 一次拷贝整个内存缓冲区

为了解决这个问题，我们应用了一个简单的优化将 cudaMemcpy()调用次数从几百次减少到一次：不再使用 Torch 框架将 Tensor 单独移动到 GPU 上，而是先将所有 Tensor 的数据放置到一块预先分配好的连续内存缓冲区中，并一次性将缓冲区内容拷贝到 GPU 里，最终通过指向 GPU 显存缓冲区的不同偏移量来重构得到 GPU 上的所有张量。

该优化带来的代价是要显式管理预先分配的内存缓冲区的生命周期，以及需要对不同数据类型手动处理 GPU 显存对齐。但作为结果，P50 数据拷贝延迟从 10ms 降低到 1ms 以下，这证明了该优化带来的复杂性是可以接受的。

4、利用 CUDA Graph

为了进一步优化模型推理过程，我们使用 CUDA Graph(https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/) 来完全消除剩余小 op 的开销。CUDA Graph 允许我们将模型推理过程捕捉为静态图，而不是单独调度。它可以让整个计算作为一个单元进行执行，而不产生任何 Kernel 启动开销。我们支持将 CUDA Graph 作为模型服务的一个新后端。一开始加载模型时，模型服务执行一次模型推理以构建图实例，该实例可以针对实时流量重复执行。

CUDA Graph 在一个批处理内（下图）执行 Kernel，而不是在一个序列中逐个执行（上图），这减少了 Kernel 之间 CPU 启动带来的开销。图表来自：https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/

CUDA Graph 自身的一些限制给我们的模型服务带来了额外的复杂性。其中最大的限制是 CUDA Graph 要求所有 Tensor 都具有静态形状和布局，这对动态大小批次和不规则的变长 Tensor 带来了挑战。然而，我们相信为了获得更好性能进行的权衡是值得的，并且我们可以将输入 Tensor 补齐到精心挑选的静态形状。

5、使用更大的 batchsize

最后，我们重新研究了 SMS 为模型推理执行的批处理策略。SMS 支持动态批处理，可以让多个请求合并成更大的批次。它通常能带来更好的吞吐量，但代价是需要较短的时间以等待请求序列收集足够多的项（item）。对于 CPU 上的 ML 推断，我们通常想要将请求分成小批量来提高并行度以减小延时。然而对于 GPU，其延时对 batchsize 并不敏感，使用更大的 batchsize 对 GPU 提升工作负载效率更重要。

这种 batchsize 的需求使我们重新思考了 SMS 里的分布式系统设置。对于 CPU 上的 ML 推断，我们使用 scatter-gather 结构将原始请求拆分，并在多个叶子结点上并行执行，以获得更小的延时。

此外，该架构允许我们为每个叶子结点分配一个固定的数据切片，以优化特征提取期间的缓存命中率。然而，由于 GPU 倾向使用大 batchsize，因此删除 root layer 直接在原始请求中使用大 batchsize 更有意义。最终我们使用了 CacheLib 的混合缓存，它用 DRAM 和 SSD 来补偿与 scatter-gather 架构设置相比的缓存容量损失。

6、结果

我们首先测量了模型单次推理的延时。我们使用 c5.18x AWS 实例提供 CPU 服务，g5.4 AWS 实例提供 GPU 服务。

CPU 延时随着 batchsize 线性增长，在较小的 batchsize 下，GPU 延时几乎相同，此时 Kernel 启动开销占延时主导因素。然而随着 batchsize 增加，实际计算时间主导了延时，GPU 延时以亚线性形式增长。在实践中，GPU 效率提升的亮点在于 SMS 可以使用更大的 batch 工作。结合了所有优化之后，我们获得了惊人的结果，相比 CPU 服务，GPU 服务在较大 batchsize 下的每批次的延迟提升了 100 倍以上。