Merlin HugeCTR 分級(jí)參數(shù)服務(wù)器簡介

參數(shù)服務(wù)器是推薦系統(tǒng)的重要組成部分，但是目前的訓(xùn)練端參數(shù)服務(wù)器由于高延遲和同步問題無法有效解決推理部署中模型過大的瓶頸。Merlin HugeCTR（以下簡稱 HugeCTR）團(tuán)隊(duì)針對(duì)傳統(tǒng)參數(shù)服務(wù)器的問題重新設(shè)計(jì)了一種分級(jí)推理端參數(shù)服務(wù)器，將 GPU 內(nèi)存作為一級(jí)緩存，Redis 集群作為二級(jí)緩存，RocksDB 作為持久化層，極大提高了推理效率。HugeCTR 團(tuán)隊(duì)將分多期為大家介紹此分級(jí)參數(shù)服務(wù)器的具體設(shè)計(jì)和細(xì)節(jié)，本期為系列的第一期。

1. 引言：

傳統(tǒng)訓(xùn)練端參數(shù)服務(wù)器及其缺陷

傳統(tǒng)參數(shù)服務(wù)器維護(hù)和同步模型參數(shù)僅用于訓(xùn)練，worker 節(jié)點(diǎn)執(zhí)行前向和后向計(jì)算。具體來說，在訓(xùn)練中：worker 節(jié)點(diǎn)從 server 節(jié)點(diǎn)中拉取其相應(yīng)的參數(shù)，進(jìn)行前向計(jì)算，通過反向傳播計(jì)算梯度，最后將這些梯度推送到服務(wù)器。在推理中，它只執(zhí)行前兩個(gè)步驟。如果部署在高性能設(shè)備集群中，worker 節(jié)點(diǎn)的計(jì)算速度非常快，因此傳統(tǒng) PS 通常會(huì)遇到這兩個(gè)瓶頸：

（1）server 和 worker 之間的 pull 和 push 操作延遲；

（2）從 worker 節(jié)點(diǎn)收到梯度后，server 節(jié)點(diǎn)中的參數(shù)同步問題。

由于 GPU 停頓、同步/一致性不足，GPU 的計(jì)算結(jié)構(gòu)很難通過使用基于 CPU 的實(shí)現(xiàn)的參數(shù)服務(wù)器來支持?jǐn)?shù)據(jù)并行。在 GPU 內(nèi)存中擬合完整模型以及小批量輸入數(shù)據(jù)和中間網(wǎng)絡(luò)狀態(tài)的需要限制了可以訓(xùn)練和推理的模型的大小。同樣的瓶頸也出現(xiàn)在推理部署中，因?yàn)橥评砉?jié)點(diǎn)也需要從集中的參數(shù)服務(wù)器組中拉取所有需要的模型參數(shù)。當(dāng)請求包含節(jié)點(diǎn)未加載的參數(shù)時(shí)，節(jié)點(diǎn)需要再次同步從參數(shù)服務(wù)器拉取參數(shù)。當(dāng)模型參數(shù)版本發(fā)生變化時(shí)，需要暫停推理服務(wù)，逐個(gè)節(jié)點(diǎn)更新參數(shù)。

使用基于 CPU 的參數(shù)服務(wù)器進(jìn)行模型推理的部署時(shí)，上述問題非常明顯，特別是推薦模型的部署。

HugeCTR 推理端分級(jí)參數(shù)服務(wù)器

為了解決上述問題，我們 HugeCTR 團(tuán)隊(duì)設(shè)計(jì)了一個(gè)全新的參數(shù)服務(wù)器系統(tǒng)，利用分布在多個(gè)服務(wù)器機(jī)器上的 GPU 擴(kuò)展深度學(xué)習(xí)應(yīng)用以進(jìn)行推理。我們稱為 HugeCTR Hierarchical GPU-based Inference Parameter Server，簡稱為 HPS。HPS 處理在并行模型實(shí)例的推理過程中使用的共享模型參數(shù)（Embedding 向量和權(quán)重）相關(guān)的同步和通信。

與其他系統(tǒng)不同，HPS 進(jìn)行了許多專門針對(duì)高效利用 GPU 的優(yōu)化，包括分布式參數(shù)服務(wù)器分片，以實(shí)現(xiàn) GPU/CPU 中許多龐大 Embedding Table 的并行推理，以及 GPU 友好的緩存、 臨時(shí)數(shù)據(jù)移動(dòng)內(nèi)存、 內(nèi)存管理機(jī)制。

推理端參數(shù)服務(wù)器的支持

支持不同模型的混合部署：如 DeepFM、DCN、DLRM、MMOE 和序列模型 （DIN、DIEN）

支持推理的大輸入數(shù)據(jù)量：Batch_size 大于 1K，look_up per request 超過 1000。

支持更快的在線熱部署：將完整模型更新/加載到推理節(jié)點(diǎn)進(jìn)行服務(wù)只需不到 10 分鐘（Embedding Table 大小大于 600G）

支持資源隔離：在推理中隔離 GPU 的內(nèi)存，確保推理服務(wù)在生產(chǎn)環(huán)境基于不同的隔離策略。例如一種支持單 GPU/CPU 的模型，通過第三方工具（如 k8s）重啟和隔離。

支持不同模型獨(dú)立的巨大 Embedding Table：不同模型每個(gè) Embedding Table 大小大于 600G。

支持單個(gè)節(jié)點(diǎn)的多級(jí)緩存：

部分參數(shù)（熱特征）可以放入 GPU 內(nèi)存直接用于推理。

部分參數(shù)（冷特征）可以放入 CPU 內(nèi)存用于 CPU 與 GPU 交換。

Embedding 的其余部分存儲(chǔ)在本地 SSD。

支持模型更新機(jī)制：

模型部署支持特定的性能指標(biāo)，在性能下降時(shí)通過版本控制及時(shí)回滾。

通過分層緩存機(jī)制來重疊數(shù)據(jù)的轉(zhuǎn)移以緩解網(wǎng)絡(luò)帶寬瓶頸。

使用最小的 PS 分片粒度不間斷更新，而無需逐節(jié)點(diǎn)更新。

支持容錯(cuò)和持久性：一個(gè)節(jié)點(diǎn)故障后參數(shù)和服務(wù)可以被恢復(fù)。

支持在線學(xué)習(xí)：每分鐘更新密集模型權(quán)重

2. HugeCTR 分級(jí)參數(shù)服務(wù)器組件塊：

CPU 分布式緩存

分布式 Redis 集群

Redis 集群的同步查找：每個(gè)模型實(shí)例從本地化的 GPU 緩存中查找所需的 Embedding key，這也會(huì)將丟失的 Embedding key（在 GPU 緩存中找不到的 key）存儲(chǔ)到丟失的鍵緩沖區(qū)中。丟失的鍵緩沖區(qū)與 Redis 實(shí)例同步交換，Redis 實(shí)例依次對(duì)任何丟失的 Embedding key 執(zhí)行查找操作。因此，分布式 Redis 集群充當(dāng)了二級(jí)緩存，可以完全替代本地化參數(shù)服務(wù)器來加載所有模型的完整 Embedding table。

GPU 緩存（Embedding Cache）

異步/同步插入

我們支持將丟失的 Embedding key 異步插入到 Embedding Cache 中。該功能可以通過配置文件中自定義的命中率閾值自動(dòng)激活。當(dāng) Embedding Cache 的真實(shí)命中率高于自定義閾值時(shí)，Embedding Cache 會(huì)異步插入缺失的 key。反之則會(huì)以同步方式插入，以確保推理請求的高精度。通過異步插入方式，與之前的同步方式相比，在 Embedding Cache 達(dá)到用戶定義的閾值后，可以進(jìn)一步提高 Embedding Cache 的真實(shí)命中率。

在線更新

我們支持將增量 Embedding Key 異步刷新到 Embedding Cache 中。當(dāng)稀疏模型文件需要更新到 GPU Embedding Cache 時(shí)，會(huì)觸發(fā)刷新操作?；谠诰€訓(xùn)練完成模型的模型版本迭代或增量參數(shù)更新后，需要將最新的 Embedding table 更新到推理服務(wù)器上的 Embedding Cache 中。為了保證運(yùn)行模型可以在線更新，我們將通過分布式事件流平臺(tái)（Kafka）更新分布式數(shù)據(jù)庫和持久化數(shù)據(jù)庫。同時(shí)，GPU Embedding Cache 會(huì)刷新現(xiàn)有 Embedding key 的值，并替換為最新的增量 Embedding vector。

本地鍵值存儲(chǔ)

本地 RocksDB 查詢引擎

對(duì)于仍然無法完全加載到 Redis 集群中的超大規(guī)模 Embdding table 我們將在每個(gè)節(jié)點(diǎn)上啟用本地鍵值存儲(chǔ)（RocksDB）。

RocksDB 的同步查詢：Redis 集群客戶端在分布式 GPU 緩存中查找 Embedding key 時(shí)，會(huì)記錄丟失的 Embedding key（在 Redis 集群中未找到的 key），記錄到丟失的 key buffer 中。丟失的 key buffer 與本地 RocksDB 客戶端同步交換，然后將嘗試在本地 SSD 中查找這些 key。最終，SSD 查詢引擎將對(duì)所有模型缺失的 Embedding key 執(zhí)行第三次查找操作。

對(duì)于已經(jīng)存儲(chǔ)在云端的模型存儲(chǔ)庫，RocksDB 將作為本地 SSD 緩存，用于存儲(chǔ) Redis 集群無法加載的剩余部分。因此，在實(shí)踐中，本地化的 RocksDB 實(shí)例充當(dāng)了三級(jí)緩存。

3. HugeCTR 分級(jí)參數(shù)服務(wù)器的配置和使用：

訓(xùn)練端

配置模型名和 Kafka broker

在訓(xùn)練段，用戶在 CreateSolver 時(shí)需要為當(dāng)前訓(xùn)練的模型提供一個(gè)模型名，這個(gè)模型名將會(huì)被參數(shù)服務(wù)器用于區(qū)分不同模型。

同時(shí)，用戶還需要配置 Kafka broker 的端口和 ip，用于將模型發(fā)送到到 Kafka。

增量模型導(dǎo)出接口

用戶可以使用以上接口，將增量模型導(dǎo)出到 Kafka broker，參數(shù)服務(wù)器端將會(huì)自動(dòng)消化 Kafka 的消息。

推理端

Embedding Cache配置：

gpucache： 用戶可自由配置是否使用 GPU 緩存。

gpucacheper： 用于決定 Embedding table 導(dǎo)入到 GPU 緩存的比例，默認(rèn)為 0.5。

hit_rate_threshold： 用戶自定義的閾值，將會(huì)決定 GPU 緩存的更新方式。

使用 hashmap/parallel hashmap 作為 CPU 緩存時(shí)的相關(guān)配置：

num_partitions： Embedding table 將會(huì)被分為多個(gè)分片進(jìn)行存儲(chǔ)，這里用于指定分片數(shù)量。

overflow_policy： 當(dāng)緩存占滿時(shí)，可選擇隨機(jī)移除或移除最老的 Embedding。

overflow_margin： 用于指定每個(gè)分片儲(chǔ)存的最大 Embedding 數(shù)量。

overflow_resolution_target： 用于指定每個(gè)分片移除 Embedding 的比例，取值為 0 到 1 之間。

initial_cache_rate： 初始的緩存率。

使用 Redis 作為 CPU 緩存時(shí)的相關(guān)配置：

需要配置服務(wù)器 ip 和端口，用戶名以及密碼，其他與 hashmap/parallel hashmap 相同。需要注意的是 num_partitions 必須大于等于 redis 節(jié)點(diǎn)的數(shù)量。

RocksDB：

path： RocksDB 存儲(chǔ)數(shù)據(jù)的路徑，由用戶自行配置。

read_only： 啟用 read_only 后，RocksDB 將無法更新，適用于靜態(tài) Embedding 的推理。

Kafka：

brokers： 用于設(shè)置 Kafka 服務(wù)器的 ip 和端口。

poll_timeout_ms： 用于數(shù)據(jù)接收的最大等待時(shí)長，超過該時(shí)間將自動(dòng)把參數(shù)更新送往數(shù)據(jù)存儲(chǔ)層。

max_receive_buffer_size： 用于數(shù)據(jù)接收緩沖區(qū)大小，超過該大小將自動(dòng)把參數(shù)更新送往數(shù)據(jù)存儲(chǔ)層。

max_batch_size： 用于設(shè)置每次批量數(shù)據(jù)發(fā)送的大小。

4. 結(jié)語

在這一期的 HugeCTR 分級(jí)參數(shù)服務(wù)器簡介中，我們介紹了傳統(tǒng)參數(shù)服務(wù)器的結(jié)構(gòu)以及 HugeCTR 分級(jí)推理參數(shù)服務(wù)器是如何在其基礎(chǔ)上進(jìn)行設(shè)計(jì)和改進(jìn)的。我們還介紹了我們的三級(jí)存儲(chǔ)結(jié)構(gòu)以及相關(guān)配置使用。在下一期中，我們將著重介紹 HugeCTR 分級(jí)參數(shù)服務(wù)器各個(gè)部件的設(shè)計(jì)細(xì)節(jié)，敬請期待。

關(guān)于作者

About Yingcan Wei

GPU計(jì)算專家，畢業(yè)于香港大學(xué)，HugeCTR算法組負(fù)責(zé)人。當(dāng)前主要從事HugeCTR的算法設(shè)計(jì)與推理架構(gòu)工作。研究領(lǐng)域包括深度學(xué)習(xí)域適應(yīng)，生成對(duì)抗網(wǎng)絡(luò)，推薦算法設(shè)計(jì)優(yōu)化。在2020年加入英偉達(dá)前，任職于歐美外資以及互聯(lián)網(wǎng)等企業(yè)，擁有多年圖像處理 、數(shù)據(jù)挖掘，推薦系統(tǒng)設(shè)計(jì)開發(fā)相關(guān)經(jīng)驗(yàn)。

About Jerry Shi

本科畢業(yè)于加州大學(xué)伯克利分校，在康奈爾大學(xué)獲得碩士文憑。于2021年加入英偉達(dá)，在Merlin HugeCTR團(tuán)隊(duì)算法組負(fù)責(zé)推薦系統(tǒng)架構(gòu)與算法的相關(guān)設(shè)計(jì)及開發(fā)。

審核編輯：郭婷