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

本實(shí)踐中，唯品會(huì) AI 平臺(tái)與 NVIDIA 團(tuán)隊(duì)合作，結(jié)合 NVIDIA TensorRT 和 NVIDIA Merlin HierarchicalKV（HKV）將推理的稠密網(wǎng)絡(luò)和熱 Embedding 全置于 GPU 上進(jìn)行加速，吞吐相比 CPU 推理服務(wù)提升高于 3 倍。

應(yīng)對(duì)GPU推理上的難題

唯品會(huì)（NYSE: VIPS）成立于 2008 年 8 月，總部設(shè)在中國(guó)廣州，旗下網(wǎng)站于同年 12 月 8 日上線。唯品會(huì)主營(yíng)業(yè)務(wù)為互聯(lián)網(wǎng)在線銷售品牌折扣商品，涵蓋名品服飾鞋包、美妝、母嬰、居家、生活等全品類。

唯品會(huì) AI 平臺(tái)服務(wù)于公司搜索、推薦、廣告等業(yè)務(wù)團(tuán)隊(duì)，提供公司級(jí)一站式服務(wù)平臺(tái)。搜索、推薦、廣告等業(yè)務(wù)旨在通過(guò)算法模型迭代，不斷優(yōu)化用戶購(gòu)買體驗(yàn)，從而提升點(diǎn)擊率和轉(zhuǎn)化率等業(yè)務(wù)指標(biāo)，最終實(shí)現(xiàn)公司銷售業(yè)績(jī)?cè)鲩L(zhǎng)。

在使用 GPU 打開(kāi)推理算力天花板過(guò)程中，遇到了如下問(wèn)題：

稠密網(wǎng)絡(luò)，如何獲取更好的 GPU 推理性能；

Embedding table 如何使用 GPU 加速查詢。

為了解決上面的問(wèn)題，我們選擇使用了 NVIDIA TensorRT 和 Merlin HierarchicalKV。具體原因如下：

稠密網(wǎng)絡(luò)使用 TensorRT 推理，通過(guò) TensorRT 和自研 Plugin 方式獲取更好的推理性能；

HierarchicalKV 是一個(gè)高性能 GPU Table 實(shí)現(xiàn)，我們將熱 Embedding 緩存在 GPU 中，冷 Embedding 則通過(guò)內(nèi)存和分布式 KV 存儲(chǔ)，加速查表過(guò)程。

GPU推理服務(wù)設(shè)計(jì)方案

AI 平臺(tái)支持搜索、推薦、廣告等所有算法業(yè)務(wù)，提供大規(guī)模分布式訓(xùn)練、推理、實(shí)時(shí)模型等基礎(chǔ)引擎平臺(tái)，打造屬于唯品會(huì)自己的 AI 基礎(chǔ)能力引擎。

圖 1. GPU 推理服務(wù)工作流程圖

如上圖所示，支持 GPU 推理服務(wù)，可以分為如下幾步：

TensorPS（自研訓(xùn)練框架）

支持離線和實(shí)時(shí)訓(xùn)練；

離線訓(xùn)練：生成天級(jí)全量模型，完成后同步給 Odin；

實(shí)時(shí)訓(xùn)練：生成小時(shí)級(jí)別的全量模型和分鐘級(jí)別的增量模型，完成后同步給 Odin；

2. Odin（模型協(xié)調(diào)者）

（離線/實(shí)時(shí)）單機(jī)模型的全量模型：觸發(fā) TensorRT Converter；

（離線/實(shí)時(shí)）分布式模型的全量模型：同時(shí)觸發(fā) TensorRT Converter 和 Reshard；

（離線/實(shí)時(shí)）單機(jī)/分布式模型的增量模型：觸發(fā) TensorRT Converter；

3. TensorRTConverter（模型轉(zhuǎn)換器）

將 Dense 網(wǎng)絡(luò)轉(zhuǎn)換成 TensorRT Engine；

轉(zhuǎn)化完成，如果是全量模型，向模型管理 API 匯報(bào)全量版本；如果是增量模型，向模型管理 API 匯報(bào)增量版本；

4. Reshard（參數(shù)分片模塊）

對(duì)模型參數(shù)分片后，向模型管理 API 匯報(bào)版本；

分片后參數(shù)，同步到分布式在線參數(shù)服務(wù) Atreus；

5. Thor（自研推理服務(wù)）

單機(jī)模型：通過(guò)模型管理 API 獲取全量模型版本，拉取模型并啟動(dòng)推理服務(wù) Thor；

分布式模型：需要部署分布式參數(shù)服務(wù) Atreus 和推理服務(wù) Thor；

如果開(kāi)啟了實(shí)時(shí)模型特性，Thor 會(huì)定時(shí)通過(guò)模型管理 API 獲取增量版本，拉取并更新增量模型；

6. Atreus（自研分布式在線參數(shù)服務(wù)）

僅用于分布式模型，可支持 TB 級(jí)參數(shù)；

如果開(kāi)啟了實(shí)時(shí)模型特性，Atreus 會(huì)定時(shí)通過(guò)模型管理 API 獲取增量版本，拉取并更新增量參數(shù)。

GPU模型推理

圖 2. 前向計(jì)算流程圖

如上圖所示，前向計(jì)算可以分為如下幾步：

H2D 拷貝（CPU -> GPU）；

Embedding 層，使用 GPU Table lookup（GPU）；

Dense 層，使用 TensorRT +自研 Plugin 推理（GPU）；

D2H 拷貝（GPU -> CPU）。

稠密網(wǎng)絡(luò)使用TensorRT在GPU上計(jì)算

圖 3. 稠密網(wǎng)絡(luò) TensorRT 推理優(yōu)化

如上圖所示：

稠密網(wǎng)絡(luò)使用 TensorRT 推理，結(jié)合自定義 Plugin 實(shí)現(xiàn)推理性能優(yōu)化。

利用HierarchicalKV實(shí)現(xiàn)GPUTablelookup

圖 4. 基于 HierarchicalKV 的 GPU Table

如上圖所示，查表過(guò)程可以分為如下幾步：

將 keys 拷貝到 GPU；

將 keys concat 成一個(gè)大的 merged keys，減少后續(xù)查表次數(shù)；

merged keys 查 GPU Table，輸出 merged values，并輸出未命中 missed keys 和 missed indices；

拷貝 missed keys 到 CPU；

查詢 Atreus（分布式參數(shù)服務(wù)器），獲取 missed values；

missed values 拷貝到 GPU；

將 missed values 更新到 merged values；

將 merged values 輸出 Split 成多個(gè) Tensor（和 keys 一一對(duì)應(yīng)）；

對(duì) missed keys 進(jìn)行去重；

去重之后，異步更新 GPU Table。

TensorRTConverter實(shí)現(xiàn)GPU模型轉(zhuǎn)換

圖 5. TensorRT Converter 轉(zhuǎn)換流程

如上圖所示，TensorRT Converter 可以分為如下幾步：

Freeze CPU 模型；

切分模型 Graph 成 Sparse 和 Dense 兩個(gè)子圖，Sparse 圖在 GPU 上執(zhí)行，Dense 圖經(jīng)過(guò)圖優(yōu)化后使用 TensorRT 推理；

Dense 圖轉(zhuǎn)化成 ONNX 模型；

優(yōu)化 ONNX 模型，把圖中 OP 替換成自定義的高性能 TensorRT Plugin；

轉(zhuǎn)換 ONNX 模型成 TensorRT Engine；

合并 Sparse 圖和 TensorRT Engine 生成 GPU 模型。

自研CUDAKernel，提高性能

GPUTable加速查表

基于 HierarchicalKV 增強(qiáng)了 find 接口，支持獲取未命中 keys indices 等信息，在高命中率情況下有更好的性能，并貢獻(xiàn)給社區(qū)：

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

2. GPU支持CSR（Compressedsparserow）格式的序列特征

根據(jù)統(tǒng)計(jì)，序列特征有 85%+的數(shù)據(jù)都是填充值，使用 CSR 格式壓縮序列特征可以大幅度減小序列特征大小。考慮到搜推序列數(shù)據(jù)的特殊性（填充值都在序列尾部），這里僅使用 value 和 offset 兩個(gè)序列表示原始稀疏矩陣，如下圖：

圖 6. CSR 的稀疏矩陣

通過(guò) Fusion 的方式，減少 Lookup 過(guò)程 CUDA Kernel 數(shù)量，提升推理性能。

圖 7. Lookup 過(guò)程優(yōu)化對(duì)比

優(yōu)化前：N 個(gè)輸入對(duì)應(yīng) N 個(gè) Lookup CUDA Kernel；

優(yōu)化后：通過(guò)提前合并，將 CUDA Kernel 數(shù)量減少為 3 個(gè)（Concat、Lookup 和 Split）。

通過(guò) Fusion 的方式，減少 CSR 處理過(guò)程 CUDA Kernel 數(shù)量，提升推理性能，下圖以 ReduceSum 舉例。

圖 8. CSR 處理優(yōu)化流程

優(yōu)化前：N 對(duì)輸入對(duì)應(yīng) N 個(gè) ReduceSum CUDA Kernel；

優(yōu)化后：通過(guò)提前合并，將 CUDA Kernel 數(shù)量減少為 4 個(gè)（2 個(gè) Concat，1 個(gè) ReduceSum 和 1 個(gè) Split）。

3.H2D，合并CPU->GPU內(nèi)存拷貝

搜推模型中有較多的特征輸入，GPU 推理中需要將這些 Tensor 從 CPU 拷貝到 GPU，頻繁小內(nèi)存的 cudaMemcpy 會(huì)導(dǎo)致性能下降，最佳實(shí)踐是將這些 Tensor 打包在一塊連續(xù)內(nèi)存中，將整個(gè)大內(nèi)存 H2D 拷貝到 GPU。

4.Tile算子融合

搜推模型中有超過(guò) 200個(gè) Tile，大量的 Kernel Launch 會(huì)帶來(lái) GPU 推理性能惡化，最佳實(shí)踐是進(jìn)行 Kernel Fusion，在一個(gè)大的算子中執(zhí)行多個(gè)小 Kernel，從而充分發(fā)揮 GPU 的并發(fā)優(yōu)勢(shì)。

持續(xù)在搜推廣場(chǎng)景中GPU加速

唯品會(huì) AI 平臺(tái)一直追求性能上的極致，未來(lái)將會(huì)持續(xù)與 NVIDIA 技術(shù)團(tuán)隊(duì)合作，繼續(xù)探索使用 HierarchicalKV 在訓(xùn)練超大型模型上的 GPU 性能優(yōu)化，在提升 GPU 性能方面進(jìn)行不斷地探索和實(shí)踐，也會(huì)對(duì) Generative Recommenders 進(jìn)行探索和實(shí)踐。

審核編輯：劉清