1. 概述

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

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

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

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

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

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

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

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

針對(duì) Persistent 后端的 Rocksdb 會(huì)以分塊查詢機(jī)制來(lái)獲取最大性能。針對(duì)不同的硬件基礎(chǔ)設(shè)施，用戶可以進(jìn)行定制化的配置，從而保證硬件資源的利用率以及推理性能的最大化。

4. 流式增量更新

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

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

5. 結(jié)語(yǔ)

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

關(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ì)開(kāi)發(fā)相關(guān)經(jīng)驗(yàn)。

About Matthias Langer

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

About Jerry Shi

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

審核編輯：郭婷