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

1. 引言：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

支持不同模型獨立的巨大 Embedding Table：不同模型每個 Embedding Table 大小大于 600G。

支持單個節(jié)點的多級緩存：

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

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

Embedding 的其余部分存儲在本地 SSD。

支持模型更新機制：

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

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

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

支持容錯和持久性：一個節(jié)點故障后參數(shù)和服務可以被恢復。

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

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

CPU 分布式緩存

分布式 Redis 集群

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

GPU 緩存（Embedding Cache）

異步/同步插入

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

在線更新

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

本地鍵值存儲

本地 RocksDB 查詢引擎

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

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

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

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

訓練端

配置模型名和 Kafka broker

在訓練段，用戶在 CreateSolver 時需要為當前訓練的模型提供一個模型名，這個模型名將會被參數(shù)服務器用于區(qū)分不同模型。

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

增量模型導出接口

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

推理端

Embedding Cache配置：

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

gpucacheper： 用于決定 Embedding table 導入到 GPU 緩存的比例，默認為 0.5。

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

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

num_partitions： Embedding table 將會被分為多個分片進行存儲，這里用于指定分片數(shù)量。

overflow_policy： 當緩存占滿時，可選擇隨機移除或移除最老的 Embedding。

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

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

initial_cache_rate： 初始的緩存率。

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

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

RocksDB：

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

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

Kafka：

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

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

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

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

4. 結(jié)語

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

關(guān)于作者

About Yingcan Wei

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

About Jerry Shi

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

審核編輯：郭婷