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

作为图片届的“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/cuC...) 以支持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/acce...) 来完全消除剩余小op的开销。CUDA Graph允许我们将模型推理过程捕捉为静态图，而不是单独调度。它可以让整个计算作为一个单元进行执行，而不产生任何Kernel启动开销。我们支持将CUDA Graph作为模型服务的一个新后端。一开始加载模型时，模型服务执行一次模型推理以构建图实例，该实例可以针对实时流量重复执行。

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

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倍以上。

我们的服务指标也展示了令人印象深刻的结果。通过优化模型操作，重新设置分布式系统以及优化数据传输并使用CUDA Graph，我们以30%的更低延时上线了77倍大的模型，并以适当成本增加了20%的吞吐量。

最后，两个数量级效率提升开启了Pinterest最先进的推荐模型架构。可以看到，模型质量显著地提升直接转化为更高的用户参与度。在过去的一年，我们以适当的基础设施成本将一个主要产品的用户参与度提升了16%。很快，我们将推出比我们的CPU模型大100倍的最大模型。

7

结论

将基于CPU模型服务转换成基于GPU服务的过程很复杂，但这是我们在Pinterest使用最先进推荐模型的一个必要步骤。我们能够以适当的成本提供大100倍的推荐模型，这为我们的机器学习工程师给Pinner解锁更相关和响应更迅速的推荐内容提供了基础。

（*本文经授权后编译发布，原文：
https://medium.com/pinterest-...*）

欢迎下载体验 OneFlow v0.8.0 最新版本： https://github.com/Oneflow-In...