人工智能/深度学习

利用 NVIDIA Merlin HierarchicalKV 实现唯品会在搜推广场景中的 GPU 推理实践

利用 NVIDIA 组件提升 GPU 推理的吞吐

本实践中,唯品会 AI 平台与 NVIDIA 团队合作,结合 NVIDIA TensorRT 和 NVIDIA Merlin HierarchicalKV(HKV)将推理的稠密网络和热 Embedding 全置于 GPU 上进行加速,吞吐相比 CPU 推理服务提升高于 3 倍。

应对 GPU 推理上的难题

唯品会(NYSE: VIPS)成立于 2008 年 8 月,总部设在中国广州,旗下网站于同年 12 月 8 日上线。唯品会主营业务为互联网在线销售品牌折扣商品,涵盖名品服饰鞋包、美妆、母婴、居家、生活等全品类。

唯品会 AI 平台服务于公司搜索、推荐、广告等业务团队,提供公司级一站式服务平台。搜索、推荐、广告等业务旨在通过算法模型迭代,不断优化用户购买体验,从而提升点击率和转化率等业务指标,最终实现公司销售业绩增长。

在使用 GPU 打开推理算力天花板过程中,遇到了如下问题:

  1. 稠密网络,如何获取更好的 GPU 推理性能;
  2. Embedding table 如何使用 GPU 加速查询。

为了解决上面的问题,我们选择使用了 NVIDIA TensorRT 和 Merlin HierarchicalKV。具体原因如下:

  1. 稠密网络使用 TensorRT 推理,通过 TensorRT 和自研 Plugin 方式获取更好的推理性能;
  2. HierarchicalKV 是一个高性能 GPU Table 实现,我们将热 Embedding 缓存在 GPU 中,冷 Embedding 则通过内存和分布式 KV 存储,加速查表过程。

GPU 推理服务设计方案

AI 平台支持搜索、推荐、广告等所有算法业务,提供大规模分布式训练、推理、实时模型等基础引擎平台,打造属于唯品会自己的 AI 基础能力引擎。

图1. GPU 推理服务工作流程图

如上图所示,支持 GPU 推理服务,可以分为如下几步:

  1. TensorPS(自研训练框架)
    1. 支持离线和实时训练;
    2. 离线训练:生成天级全量模型,完成后同步给 Odin;
    3. 实时训练:生成小时级别的全量模型和分钟级别的增量模型,完成后同步给 Odin;
  2. Odin(模型协调者)
    1. (离线/实时)单机模型的全量模型:触发 TensorRT Converter;
    2. (离线/实时)分布式模型的全量模型:同时触发 TensorRT Converter 和 Reshard;
    3. (离线/实时)单机/分布式模型的增量模型:触发 TensorRT Converter;
  3. TensorRT Converter(模型转换器)
    1. 将 Dense 网络转换成 TensorRT Engine;
    2. 转化完成,如果是全量模型,向模型管理 API 汇报全量版本;如果是增量模型,向模型管理 API 汇报增量版本;
  4. Reshard(参数分片模块)
    1. 对模型参数分片后,向模型管理 API 汇报版本;
    2. 分片后参数,同步到分布式在线参数服务 Atreus;
  5. Thor(自研推理服务)
    1. 单机模型:通过模型管理 API 获取全量模型版本,拉取模型并启动推理服务 Thor;
    2. 分布式模型:需要部署分布式参数服务 Atreus 和 推理服务 Thor;
    3. 如果开启了实时模型特性,Thor 会定时通过模型管理 API 获取增量版本,拉取并更新增量模型;
  6. Atreus(自研分布式在线参数服务)
    1. 仅用于分布式模型,可支持 TB 级参数;
    2. 如果开启了实时模型特性,Atreus 会定时通过模型管理 API 获取增量版本,拉取并更新增量参数。

GPU 模型推理

图 2. 前向计算流程图

如上图所示,前向计算可以分为如下几步:

  1. H2D 拷贝(CPU -> GPU);
  2. Embedding 层,使用 GPU Table lookup(GPU);
  3. Dense 层,使用 TensorRT + 自研 Plugin 推理(GPU);
  4. D2H 拷贝(GPU -> CPU)。

稠密网络使用 TensorRT 在 GPU 上计算

图 3. 稠密网络 TensorRT 推理优化

如上图所示:

  1. 稠密网络使用 TensorRT 推理,结合自定义 Plugin 实现推理性能优化。

利用 HierarchicalKV 实现 GPU Table lookup

图 4. 基于 HierarchicalKV 的 GPU Table

如上图所示,查表过程可以分为如下几步:

  1. 将 keys 拷贝到 GPU;
  2. 将 keys concat 成一个大的 merged keys,减少后续查表次数;
  3. merged keys 查 GPU Table,输出 merged values,并输出未命中 missed keys 和 missed indices;
  4. 拷贝 missed keys 到 CPU;
  5. 查询 Atreus(分布式参数服务器),获取 missed values;
  6. missed values 拷贝到 GPU;
  7. 将 missed values 更新到 merged values;
  8. 将 merged values 输出 Split 成多个 Tensor(和 keys 一一对应);
  9. 对 missed keys 进行去重;
  10.  去重之后,异步更新 GPU Table。

TensorRT Converter 实现 GPU 模型转换

图 5. TensorRT Converter 转换流程

如上图所示,TensorRT Converter 可以分为如下几步:

  1. Freeze CPU 模型;
  2. 切分模型 Graph 成 Sparse 和 Dense 两个子图,Sparse 图在 GPU 上执行,Dense 图经过图优化后使用 TensorRT 推理;
  3. Dense 图转化成 ONNX 模型;
  4. 优化 ONNX 模型,把图中 OP 替换成自定义的高性能 TensorRT Plugin;
  5. 转换 ONNX 模型成 TensorRT Engine;
  6. 合并 Sparse 图和 TensorRT Engine 生成 GPU 模型。

自研 CUDA Kernel,提高性能

  1. GPU Table 加速查表

基于 HierarchicalKV 增强了 find 接口,支持获取未命中 keys indices 等信息,在高命中率情况下有更好的性能,并贡献给社区:

void find(const size_type n, 
           const key_type* keys,                     // (n)
           value_type* values,                       // (n, DIM)
           key_type* missed_keys,                    // (n)
           int* missed_indices,                      // (n)
           int* missed_size,                         // scalar
           score_type* scores = nullptr,             // (n)
           cudaStream_t stream = 0) const
  1. GPU 支持 CSR(Compressed sparse row)  格式的序列特征
  2. 根据统计,序列特征有 85%+ 的数据都是填充值,使用 CSR 格式压缩序列特征可以大幅度减小序列特征大小。考虑到搜推序列数据的特殊性(填充值都在序列尾部),这里仅使用 value 和 offset 两个序列表示原始稀疏矩阵,如下图:

图 6. CSR 的稀疏矩阵

  • 通过 Fusion 的方式,减少 Lookup 过程 CUDA Kernel 数量,提升推理性能。

图 7. Lookup 过程优化对比

  • 优化前:N 个输入对应 N 个 Lookup CUDA Kernel;
    • 优化后:通过提前合并,将 CUDA Kernel 数量减少为 3 个(Concat、Lookup 和 Split)。
  • 通过 Fusion 的方式,减少 CSR 处理过程 CUDA Kernel 数量,提升推理性能,下图以 ReduceSum 举例。

图 8. CSR 处理优化流程

  • 优化前:N 对输入对应 N 个 ReduceSum CUDA Kernel;
    • 优化后:通过提前合并,将 CUDA Kernel 数量减少为 4 个(2 个 Concat,1 个 ReduceSum 和 1 个 Split)。
  • H2D,合并 CPU -> GPU 内存拷贝

搜推模型中有较多的特征输入,GPU 推理中需要将这些 Tensor 从 CPU 拷贝到 GPU,频繁小内存的 cudaMemcpy 会导致性能下降,最佳实践是将这些 Tensor 打包在一块连续内存中,将整个大内存 H2D 拷贝到 GPU。 

  • Tile 算子融合

搜推模型中有超过 200 个 Tile,大量的 Kernel Launch 会带来 GPU 推理性能恶化,最佳实践是进行 Kernel Fusion,在一个大的算子中执行多个小 Kernel,从而充分发挥 GPU 的并发优势。

持续在搜广推场景中 GPU 加速

唯品会 AI 平台,一直追求性能上的极致,未来将会持续与 NVIDIA 技术团队合作,继续探索使用 HierarchicalKV 在训练超大型模型上的 GPU 性能优化,在提升 GPU 性能方面进行不断地探索和实践,也会对 Generative Recommenders 进行探索和实践。

标签