Merlin HugeCTR分級參數(shù)服務器：緩存和在線更新設計

1. 概述

HPS 數(shù)據(jù)后端作為 GPU embedding 緩存架構的基石，同時也是 GPU embedding 緩存在 CPU 內(nèi)存以及本地磁盤的進一步物理擴展。HPS 數(shù)據(jù)后端通過綁定不同物理層級的存儲從而提供了大型模型 embedding table 的緩存，查詢，更新以及容錯等服務，目的即為了保證在推理服務中 GPU embedding 緩存的高命中率，從而提高推理服務的吞吐大幅度降低端到端的延遲。

2. Volatile 數(shù)據(jù)后端

Volatile 數(shù)據(jù)后端以 RAM Memory 作為主要存儲介質(zhì)，提供本地或者遠端更加快捷的參數(shù)讀寫服務。既可以作為 GPU embedding 緩存的擴展，也可以作為本地磁盤的（Persistent 后端）緩存。由于 Persistent 數(shù)據(jù)后端的存儲容量可以理解為無限巨大，但是同時也就意味著在實際的推理服務中的讀寫速度也是相對較慢的。因此 Volatile 數(shù)據(jù)后端彌補了 Persistent 數(shù)據(jù)庫的缺陷，極大擴展了有限的 GPU 內(nèi)存。但是 Volatile 數(shù)據(jù)庫的容量本質(zhì)上也是有限的，因此我們通過實現(xiàn)參數(shù)的分區(qū)，驅(qū)逐以及共享進一步提高 RAM Memory 帶來的上限。

為了保證 Volatile 數(shù)據(jù)后端可以適用于更廣泛的推薦部署場景，我們同樣實現(xiàn)了針對不同場景的本地化或者分布式存儲機制。從下圖中我們可以看到，針對 HashMap 此類常規(guī)的存儲結構，通過實例化 Volatile 數(shù)據(jù)后端，可以保證每個數(shù)據(jù)節(jié)點都將擁有一個獨立本地化的存儲實例。對于本地化的 Volatile 數(shù)據(jù)后端，我們不僅實現(xiàn)了分區(qū)優(yōu)化的 HashMap 結構，還實現(xiàn)了高性能的 Parallel HashMap 數(shù)據(jù)結構，進而大幅度提升了本地參數(shù)的查詢和讀寫服務。

與之相反的，則是分布式共享 Volatile 數(shù)據(jù)后端，分布式的數(shù)據(jù)后端將參數(shù)通過邏輯分區(qū)保存在不同的網(wǎng)絡節(jié)點，即 Redis Cluster（既可以將參數(shù)分布式存儲在遠端的數(shù)據(jù)節(jié)點，同樣也可以是相同的推理節(jié)點）。通過使用集群中共享 RAM 內(nèi)存進行參數(shù)的存儲和讀寫，進一步擴大了 Volatile 數(shù)據(jù)庫容量，也提供了 Redis 持久化特性（RDB 和 AOF 等 ），從而使得跨節(jié)點重啟之后的讀寫操作可以無縫進行。 由于共享的機制，也實現(xiàn)了 HugeCTR 模型訓練到推理的參數(shù)無縫更新。

為了最大限度地提高性能并避免由零星 RAM memory 在 Volatile 數(shù)據(jù)后端中的高效使用（即內(nèi)存不足的情況），我們提供了溢出處理機制。 它允許限制每個分區(qū)存儲的最大嵌入量，從而限制分布式數(shù)據(jù)庫的內(nèi)存消耗。當前我們允許用戶配置不同的驅(qū)逐更新策略（隨機驅(qū)逐以及 LRU）來保證內(nèi)存最大限度的利用。對于本地化的參數(shù)分區(qū)以及查詢機制，用戶可以對數(shù)據(jù)后端顯式的配置分區(qū)數(shù)量以及分區(qū)大小來更加細粒度的提高數(shù)據(jù)后端的查詢讀寫服務。具體來說，用戶可以通過配置 max_get_batch_size，max_set_batch_size 以及并發(fā)線程數(shù)來精準控制讀寫開銷。

具體參數(shù)配置信息可以參考 HPS 配置詳解：https://github.com/triton-inference-server/hugectr_backend/blob/main/docs/hierarchical_parameter_server.md#6-configuration

3. Persistent 數(shù)據(jù)后端

Persistent 數(shù)據(jù)后端以本地 SSDs 作為持久化存儲介質(zhì)，維護一個完整的模型參數(shù)部分，同時承載著模型參數(shù)的容錯的功能。相對于 GPU embedding 緩存與 Volatile RAM 緩存，Persistent 數(shù)據(jù)后端可以看作是一個擁有無限虛擬空間的存儲后端，同時也作為一個本地化的 Key-value 查詢引擎，在此我們引入了性能優(yōu)化后的 RocksDB 作為 Persistent 后端的實現(xiàn)。

每個推理節(jié)點通過 HPS 的配置文件即可在本地磁盤保留所有模型 embedding table 的獨立完整副本。 Persistent 數(shù)據(jù)后端也是對分布式 Volatile 數(shù)據(jù)后端的進一步補充：1） 進一步擴展存儲容量的同時， 2） 實現(xiàn)高可用性。 特別是對于超大規(guī)模的模型（甚至超過了分布 Redis Cluster 的總 RAM 容量），或者由于網(wǎng)絡帶寬等硬件限制造成 Redis Cluster 不可用，RocksDB 的分成存儲結構 同樣可以完全滿足高并發(fā)參數(shù)查詢請求。

針對 Persistent 后端的 Rocksdb 會以分塊查詢機制來獲取最大性能。針對不同的硬件基礎設施，用戶可以進行定制化的配置，從而保證硬件資源的利用率以及推理性能的最大化。

4. 流式增量更新

通過優(yōu)化后 Kafka 的發(fā)布更新機制，推理節(jié)點中的每個 HPS 實例所對應的不同層級的數(shù)據(jù)后端通過訂閱對應模型的 Topic，消費實時的增量模型參數(shù)，實現(xiàn)模型的異步更新。在具體實現(xiàn)中，我們提供兩個簡單易用的抽象接口，分別是 MessageSink 和 MessageSource，保證了增量模型從訓練端向推理端的無縫更新。

用戶可以通過 JSON 格式的配置文件，任意組合搭配適用于部署場景的數(shù)據(jù)后端，保證每個部署節(jié)點充分的利用上文所提到的所有存儲介質(zhì)。由于 HPS 是構建于 Triton 推理架構之上。正如下圖所展示的完整 HPS 訓練推理數(shù)據(jù)流示意圖，每個 Triton 節(jié)點通過 Volatile 數(shù)據(jù)后端既可以維護高性能的本地 RAM 緩存，同時也可以負責維護對應 Redis 節(jié)點的參數(shù)分區(qū)，Redis cluster 中的分區(qū)參數(shù)既可以通過訓練與推理集群共享，也可以通過訂閱 Kafka 數(shù)據(jù)源來實現(xiàn)無縫的參數(shù)更新。模型在 HugeCTR 訓練平臺可以實時將在線訓練中的增量模型推向分布式的 Kafka 隊列，相同 JSON 配置的推理節(jié)點會自動檢測和監(jiān)控消息隊列，從而保證了推理節(jié)點的 Persistent 數(shù)據(jù)后端始終維護完整的最新版本的參數(shù)副本，為推理服務的容錯提供了保障。

5. 結語

在這一期的 HugeCTR 分級參數(shù)服務器介紹中，我們介紹了 CPU 分布式緩存，本地緩存，以及在線更新的設計細節(jié)。在下一期中，我們將著重介紹 HugeCTR 分級參數(shù)服務器中最關鍵的組件：Embedding Cache 的設計細節(jié)，敬請期待。

關于作者：

About Yingcan Wei

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

About Matthias Langer

Matthias 在西門子擔任工廠自動化軟件和機器學習方面的調(diào)試專家多年。 2019 年，他在拉籌伯大學（墨爾本）獲得計算機科學博士學位，以表彰他在帶寬受限環(huán)境中的分布式深度學習方面的工作。 Matthias 與澳大利亞體育學院 （AIS） 為東京奧運會開發(fā)了一種用于視頻分析的新型 AI 系統(tǒng)，并開發(fā)了算法來幫助澳大利亞移動電話網(wǎng)絡中的負載預測。此外，他曾在中國領先的技術招聘平臺 BOSS 直聘北京研究中心擔任機器科學家，直至 2021 年 6 月首次公開募股。 在 NVIDIA，Matthias 是上海 DevTech APAC 團隊的一員，該團隊負責開發(fā) NVIDIA HugeCTR。 Matthias 的工作重點是分發(fā)、并行化和性能優(yōu)化算法，這些算法使 HugeCTR 成為世界上用于推薦模型訓練和推理的最快和最具可擴展性的系統(tǒng)之一。目前主要負責 HugeCTR 的分布式嵌入存儲和檢索模塊的實現(xiàn)、維護和擴展。

About Jerry Shi

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

審核編輯：郭婷