利用 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 進行探索和實踐。
