利用 NVIDIA 组件提升 GPU 推理的吞吐
本实践中，唯品会 AI 平台与 NVIDIA 团队合作，结合 NVIDIA TensorRT 和 NVIDIA Merlin HierarchicalKV（HKV）将推理的稠密网络和热 Embedding 全置于 GPU 上进行加速，吞吐相比 CPU 推理服务提升高于 3 倍。
应对 GPU 推理上的难题
唯品会（NYSE: VIPS）成立于 2008 年 8 月，总部设在中国广州，旗下网站于同年 12 月 8 日上线。唯品会主营业务为互联网在线销售品牌折扣商品，涵盖名品服饰鞋包、美妆、母婴、居家、生活等全品类。
唯品会 AI 平台服务于公司搜索、推荐、广告等业务团队，提供公司级一站式服务平台。搜索、推荐、广告等业务旨在通过算法模型迭代，不断优化用户购买体验，从而提升点击率和转化率等业务指标，最终实现公司销售业绩增长。
在使用 GPU 打开推理算力天花板过程中，遇到了如下问题：
- 稠密网络，如何获取更好的 GPU 推理性能；
- Embedding table 如何使用 GPU 加速查询。
为了解决上面的问题，我们选择使用了 NVIDIA TensorRT 和 Merlin HierarchicalKV。具体原因如下：
- 稠密网络使用 TensorRT 推理，通过 TensorRT 和自研 Plugin 方式获取更好的推理性能；
- HierarchicalKV 是一个高性能 GPU Table 实现，我们将热 Embedding 缓存在 GPU 中，冷 Embedding 则通过内存和分布式 KV 存储，加速查表过程。
GPU 推理服务设计方案
AI 平台支持搜索、推荐、广告等所有算法业务，提供大规模分布式训练、推理、实时模型等基础引擎平台，打造属于唯品会自己的 AI 基础能力引擎。
图1. GPU 推理服务工作流程图
如上图所示，支持 GPU 推理服务，可以分为如下几步：
- TensorPS（自研训练框架）
- 支持离线和实时训练；
- 离线训练：生成天级全量模型，完成后同步给 Odin；
- 实时训练：生成小时级别的全量模型和分钟级别的增量模型，完成后同步给 Odin；
- Odin（模型协调者）
- （离线/实时）单机模型的全量模型：触发 TensorRT Converter；
- （离线/实时）分布式模型的全量模型：同时触发 TensorRT Converter 和 Reshard；
- （离线/实时）单机/分布式模型的增量模型：触发 TensorRT Converter；
- TensorRT Converter（模型转换器）
- 将 Dense 网络转换成 TensorRT Engine；
- 转化完成，如果是全量模型，向模型管理 API 汇报全量版本；如果是增量模型，向模型管理 API 汇报增量版本；
- Reshard（参数分片模块）
- 对模型参数分片后，向模型管理 API 汇报版本；
- 分片后参数，同步到分布式在线参数服务 Atreus；
- Thor（自研推理服务）
- 单机模型：通过模型管理 API 获取全量模型版本，拉取模型并启动推理服务 Thor；
- 分布式模型：需要部署分布式参数服务 Atreus 和 推理服务 Thor；
- 如果开启了实时模型特性，Thor 会定时通过模型管理 API 获取增量版本，拉取并更新增量模型；
- Atreus（自研分布式在线参数服务）
- 仅用于分布式模型，可支持 TB 级参数；
- 如果开启了实时模型特性，Atreus 会定时通过模型管理 API 获取增量版本，拉取并更新增量参数。
GPU 模型推理
图 2. 前向计算流程图
如上图所示，前向计算可以分为如下几步：
- H2D 拷贝（CPU -> GPU）；
- Embedding 层，使用 GPU Table lookup（GPU）；
- Dense 层，使用 TensorRT + 自研 Plugin 推理（GPU）；
- D2H 拷贝（GPU -> CPU）。
稠密网络使用 TensorRT 在 GPU 上计算
图 3. 稠密网络 TensorRT 推理优化
如上图所示：
- 稠密网络使用 TensorRT 推理，结合自定义 Plugin 实现推理性能优化。
利用 HierarchicalKV 实现 GPU Table lookup
图 4. 基于 HierarchicalKV 的 GPU Table
如上图所示，查表过程可以分为如下几步：
- 将 keys 拷贝到 GPU；
- 将 keys concat 成一个大的 merged keys，减少后续查表次数；
- merged keys 查 GPU Table，输出 merged values，并输出未命中 missed keys 和 missed indices；
- 拷贝 missed keys 到 CPU；
- 查询 Atreus（分布式参数服务器），获取 missed values；
- missed values 拷贝到 GPU；
- 将 missed values 更新到 merged values；
- 将 merged values 输出 Split 成多个 Tensor（和 keys 一一对应）；
- 对 missed keys 进行去重；
- 去重之后，异步更新 GPU Table。
TensorRT Converter 实现 GPU 模型转换
图 5. TensorRT Converter 转换流程
如上图所示，TensorRT Converter 可以分为如下几步：
- Freeze CPU 模型；
- 切分模型 Graph 成 Sparse 和 Dense 两个子图，Sparse 图在 GPU 上执行，Dense 图经过图优化后使用 TensorRT 推理；
- Dense 图转化成 ONNX 模型；
- 优化 ONNX 模型，把图中 OP 替换成自定义的高性能 TensorRT Plugin；
- 转换 ONNX 模型成 TensorRT Engine；
- 合并 Sparse 图和 TensorRT Engine 生成 GPU 模型。
自研 CUDA Kernel，提高性能
- 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- GPU 支持 CSR（Compressed sparse row） 格式的序列特征
- 根据统计，序列特征有 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 进行探索和实践。
