[源碼解析] NVIDIA HugeCTR，GPU版本參數服務器 --(1)

0x00 摘要

本系列之中我們將會介紹 NVIDIA 出品的 HugeCTR，這是一個面向行業的推薦系統訓練框架，針對具有模型並行嵌入和數據並行密集網絡的大規模 CTR 模型進行了優化。

本文以Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems，GitHub 源碼文檔 https://github.com/NVIDIA-Merlin/HugeCTR 的翻譯爲基礎，並且結合源碼進行分析。

其中借鑑了HugeCTR源碼閱讀這篇大作，特此感謝，期望能在此篇大作基礎之上，再豐富一下對HugeCTR的理解。

0x01 背景

我們將簡要討論 CTR 估計在現代推薦系統中的作用及其訓練中的主要挑戰。

1.1 推薦系統中的點擊率估計

從在線廣告和電子商務到流媒體服務，推薦系統無處不在，同時對服務提供商的收入產生巨大影響。推薦系統找到給定用戶最可點擊的項目，然後對它們進行排名並向用戶顯示前 N 個項目。爲了實現這個目標，推薦系統首先必須估計特定用戶點擊項目的可能性。此任務通常稱爲 CTR 估計。

如何估算點擊率？這裏沒有巫術，一般是獲取包含用戶-物品交互的富數據集，並使用它來訓練 ML 模型。數據集中的每條記錄都可以包含來自用戶（年齡、工作），商品（類型、價格）和用戶商品點擊（0 或 1）的特徵。例如，如果用戶 A 從一系列書籍中購買或點擊了幾本傳記，那麼模型爲傳記分配高概率值是有意義的。

CTR 的系統結構大致如下：

下圖展示了CTR推理流程。

圖來自HugeCTR_Webinar

1.2 點擊率估算訓練的挑戰

首先，推薦系統之中的特徵有如下性質：高維，稀疏。大規模推薦系統會面臨用戶和物品的頻繁變化，因此識別用戶點擊背後的隱式特徵交互至關重要，這樣推薦系統可以提供更高質量的更通用的推薦。例如，30 歲以下的已婚人士和孩子未滿 2 歲的人可能傾向於購買高 ABV 的啤酒。對這些隱式特徵交互進行建模需要領域專家進行復雜的特徵工程。更糟糕的是，由於特徵極其複雜且不直觀，即使是人類專家也常常無法發現這些交互。爲了代替這種對專家的依賴，人們研究出了一些基於深度學習的方法，例如 Wide & Deep，DeepFM 和 DLRM，這些模型可以捕獲這些複雜的交互。

訓練 CTR 估計模型的另一個挑戰是用戶和物品幾乎每天都在變化，因此訓練出來的模型其生命週期可能很短。此外，由於數據集的大小的增加，維數和稀疏性因素，CTR 模型通常包含一個很大的嵌入表，其可能無法放入單個 GPU 甚至多個 GPU 的節點中。因此，數據加載，嵌入表查找和 GPU 間通信可以佔據模型訓練時間的很大一部分。

這些因素，再加上缺乏用於 CTR 估算的標準化建模方法，通常導致服務在吞吐量和延遲方面經常只能達到次優性能。所以在單個或多個 GPU 上完成模型的更快迭代訓練是非常重要的。

0x02 HugeCtr

HugeCTR 是一個開源框架，用於在 NVIDIA GPU 上加速 CTR 估計模型的訓練，並針對 NVIDIA GPU 的性能進行了高度優化，同時允許用戶以 JSON 格式自定義模型。它是用 CUDA C++ 編寫的，並且高度利用了 GPU 加速庫，例如cuBLAS、cuDNN和NCCL。它最初是作爲內部原型來評估 GPU 在 CTR 估計問題上的潛力，但是其很快成爲基於 GPU 的推薦系統的參考設計。由於它自然而然地成爲了專用於 CTR 估算的更通用的框架，因此 NVIDIA 於 2019 年 9 月開源了其初始版本，以接受外部反饋，同時與一些客戶保持互動。

HugeCTR 也是 NVIDIA Merlin的支柱，這是一個框架和生態系統，用於構建需要大量數據集進行訓練的大規模推薦系統，旨在促進推薦系統開發的所有階段，並在 NVIDIA GPU 上加速。

圖來自源碼 https://github.com/NVIDIA-Merlin/Merlin

HugeCTR 在單個 NVIDIA V100 GPU 上的速度比 TensorFlow 在 40 核 CPU 節點上提高了 114 倍，在同一個 V100 GPU 上實現了 TensorFlow 的 8.3 倍提高。由於由線性模型和深度模型組成的混合模型已變得普遍，因此 HugeCTR 架構 2.1 版擴展爲支持 Wide & Deep、DCN 和 DeepFM 等模型。更新包括新的數據讀取器，它可以同時讀取連續和分類輸入數據；以及新的層，包括因子分解機和交叉層。爲了實現更靈活的設計空間探索，還添加了 Dropout、L1/L2 正則化器等。

0x03 架構

3.1 CTR DL 模型

下圖描繪了用於 CTR 估計的 DL 模型的步驟：

按批次讀取數據記錄，每個記錄都由高維、極其稀疏（或 categorical 類型）的特徵組成。每個記錄還可以包含密集的數字特徵，這些特徵可以直接饋送到全連接層。
使用嵌入層將輸入稀疏特徵壓縮爲低維密集嵌入向量。例如，如果有 N 個稀疏特徵，嵌入維度爲 K，則嵌入表生成 N 個 K 維密集向量。
使用前饋神經網絡來估計點擊率。

圖上顯示了一個典型的 CTR 模型，包括數據讀取器、嵌入和全連接層。圖來自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

3.2 HugeCTR 架構

HugeCTR 不僅支持 CTR DL 所有三個步驟，而且還增強了端到端的性能，比如：

爲了防止數據加載成爲訓練中的主要瓶頸，它實現了一個專用的數據讀取器，該讀取器是異步和多線程的。它將讀取一組批處理數據記錄，其中每條記錄都由高維、極度稀疏或分類特徵（categorical features）組成。每個記錄還可以包含密集的數字特徵（dense numerical features），這些特徵可以直接饋送到全連接層。
嵌入層用於將稀疏輸入特徵壓縮爲低維、密集的嵌入向量。共有三個 GPU 加速的嵌入階段：
- 表查找
- 每個插槽（slot）內的權重規約。
- 跨插槽的權重拼接（concatenation）。
通過利用高效的 CUDA 優化技術和支持 CUDA 的庫來支持前向和後向傳播中的所有層，優化器和損失函數都是在 CUDA C++ 中實現的。

爲了訓練大規模 CTR 估計模型，HugeCTR 中的嵌入表是模型並行的，並分佈在同構集羣中的所有 GPU 上，該集羣由多個節點組成。每個 GPU 都有自己的：

前饋神經網絡（數據並行）來估計點擊率。
哈希表使數據預處理更容易並啓用動態插入。

所以，可以擴展到多個 GPU 和節點的HugtCTR的架構總結如下：

3.3 基於GPU的參數服務器

HugeCTR 實現的是一個基於GPU的參數服務器，其將embedding層放到GPU之中，worker通過與參數服務器的交互來獲取embedding。

圖來自HugeCTR_Webinar

0x04 核心功能

在本節中，我們將介紹 HugeCTR 的關鍵特性，這些特性有助於其高性能和可用性。注意：多節點訓練和混合精度訓練可以同時使用。

4.1 模型並行訓練

HugeCTR 原生支持模型並行和數據並行訓練，使得在 GPU 上訓練非常大的模型成爲可能。

4.1.1 in-memory GPU hash table

在 CTR 估計中，嵌入（embedding）對於獲得不錯的模型精度幾乎是必不可少的。它通常會導致對內存容量和帶寬的高需求以及相當數量的並行性。如果embedding分佈在多個 GPU 或多個節點上，則通信開銷也可能很大。由於用戶和物品數量龐大且不斷增加，龐大的嵌入表在所難免。

爲了克服這些挑戰並實現更快的訓練，HugeCTR實現了自己的嵌入層，其中包括一個 GPU 加速的哈希表，並利用NCCL 作爲其 GPU 間通信原語。哈希表的實現基於RAPIDS cuDF 的實現，RAPIDS cuDF 是來自 NVIDIA 的 GPU DataFrame 庫。cuDF GPU 哈希表可以比 Threading Building Blocks (TBB) 的 concurrent_hash_map 多出高達 35 倍的加速。

總之，HugeCTR 支持跨越同構計算集羣中的多個 GPU 和多個節點的模型並行嵌入表。嵌入的特徵和類別可以分佈在多個 GPU 和節點上。例如，如果您有兩個具有 8xA100 80GB GPU 的節點，則可以完全在 GPU 上訓練大至 1TB 的模型。通過使用嵌入訓練緩存，您可以在相同節點上訓練更大的模型。

4.1.2 Multi-slot embedding

嵌入表可以被分割成多個槽（或feature fields）。在嵌入查找過程中，屬於同一槽的稀疏特徵輸入在分別轉換爲相應的密集嵌入向量後，被簡化爲單個嵌入向量。然後，來自不同槽的嵌入向量連接在一起。

多槽（multi-slot）嵌入通過以下方式提高了 GPU 間帶寬利用率：

當數據集中有很多特徵時，它有助於將每個槽中有效特徵的數量減少到可管理的程度。
通過拼接不同插槽的輸出，它減少了 GPU 之間的事務數量，從而促進了更高效的通信。

下圖顯示了操作序列和 GPU 間通信 ( all2all) 是如何發生的。

該圖顯示了一個跨越 4 個 GPU 的模型並行嵌入，以及它如何與這些 GPU 的神經網絡進行交互。它還顯示瞭如何減少每個插槽的輸入特徵並跨兩個插槽連接。圖來自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

多槽嵌入對線性模型也很有用，它基本上是特徵的加權和，只需將槽數和嵌入維度都設置爲 1 即可。有關更多信息，請參閱Wide & Deep 示例。

4.1.3 具體實現

爲了在不同的嵌入上獲得最佳性能，可以選擇不同的嵌入層實現。這些實現中的每一個都針對不同的實際培訓案例，例如：

LocalizedSlotEmbeddingHash：同一個槽（特徵域）中的特徵會存儲在一個GPU中，這就是爲什麼它被稱爲“本地化槽”，根據槽的索引號，不同的槽可能存儲在不同的GPU中。LocalizedSlotEmbedding 針對每個embedding 小於 GPU 內存大小的實例進行了優化。由於在 LocalizedSlotEmbedding 中使用了每個插槽的局部規約（查完 embedding 得到向量之後，因爲已經拿到了這個slot 的所有 embedding，可以做完pooling之後再做多GPU卡通信），而在 GPU 之間沒有全局規約，因此 LocalizedSlotEmbedding 中的整體數據傳輸量遠小於 DistributedSlotEmbedding。

注意：確保輸入數據集中沒有任何重複的鍵。
DistributedSlotEmbeddingHash：所有特徵都存儲於不同特徵域/槽上，不管槽索引號是多少，這些特徵都根據特徵的索引號分佈到不同的GPU上。這意味着同一插槽中的特徵可能存儲在不同的 GPU 中，這就是將其稱爲“分佈式插槽”的原因。由於需要全局規約，所以 DistributedSlotEmbedding 適合 embedding 大於 GPU 內存大小的情況，因而 DistributedSlotEmbedding 在 GPU 之間有更多的內存交換。

注意：確保輸入數據集中沒有任何重複的鍵。
LocalizedSlotEmbeddingOneHot：一種特殊的 LocalizedSlotEmbedding，需要一個獨熱數據輸入。每個特徵字段也必須從零開始索引。例如，性別應該是0,1，而1,2 就不正確。

一定要注意，LocalizedSlotEmbeddingHash 和 DistributedSlotEmbeddingHash 的區別在於同一個槽（特徵域）中的特徵 是不是 會存儲在同一個GPU中。比如，有 2 張GPU卡，有4個slot。

local 模式：GPU0 存 slot0 和 slot1，GPU1 存 slot2 和 slot3。
distribute 模式：每個 GPU 都會存所有 slot 的一部分參數，通過哈希方法決定如何將一個參數分配到哪個 GPU 上。

4.2 多節點訓練

多節點訓練使得我們很容易訓練任意大小的嵌入表。在多節點解決方案中，稀疏模型（稱爲嵌入層）分佈在節點之間。同時，密集模型（例如 DNN）是數據並行的，並且在每個 GPU 中都包含密集模型的副本（見下圖）。通過我們的實施，HugeCTR 利用 NCCL 進行高速和可擴展的節點間和節點內通信。

圖來自源碼。

要在多個節點上運行，HugeCTR 應該使用 OpenMPI 構建。建議支持GPUDirect RDMA以獲得高性能。有關更多信息，請參閱DCN 多節點訓練樣本。

4.3 混合精度訓練

混合精度訓練已成爲在保持模型精度的同時實現進一步加速的常用技術，可以幫助我們改善和減少內存吞吐量佔用。在 HugeCTR 中，可以配置全連接層以利用 NVIDIA Volta 架構及其後續架構上的張量核心。它們在內部使用 FP16 進行加速矩陣乘法，但其輸入和輸出仍爲 FP32。

混合精度訓練在這種模式下，TensorCores 被用於提高基於矩陣乘法的層的性能，例如FullyConnectedLayer和InteractionLayer，在 Volta、Turing 和 Ampere 架構上。對於包括嵌入在內的其他層，數據類型更改爲 FP16，以便節省內存帶寬和容量。要啓用混合精度模式，請在配置文件中指定 mix_precision 選項。當mixed_precision設定，完整的FP16管道將被觸發。將應用損失縮放以避免算術下溢（見圖）。可以使用配置文件啓用混合精度訓練。

圖 5：算術下溢圖來自源碼。

4.4 SGD 優化器和學習率調度

學習率調度允許用戶配置其超參數，包括以下內容：

learning_rate：基礎學習率。
warmup_steps：用於預熱的初始步驟數。
decay_start：指定學習率衰減開始的時間。
decay_steps：衰減期（逐步）。

圖 6 說明了這些超參數如何與實際學習率相互作用。

有關更多信息，請參閱Python 接口。

圖 6：學習率調度圖來自源碼。

4.5 嵌入訓練緩存

嵌入訓練緩存（Model Oversubscription）使您能夠訓練高達 TB 的大型模型。它是通過在訓練階段以粗粒度、按需方式將超過 GPU 內存聚合容量的嵌入表的一個子集加載到 GPU 中來實現的。要使用此功能，您需要將數據集拆分爲多個子數據集，同時從中提取唯一鍵集（見圖 7）。

此功能目前支持單節點和多節點訓練。它支持所有嵌入類型，並且可以與Norm和Raw數據集格式一起使用。我們修改了我們的criteo2hugectr工具以支持 Criteo 數據集的密鑰集提取。有關更多信息，請參閱我們的Python Jupyter Notebook，瞭解如何將此功能與 Criteo 數據集結合使用。

注意：Criteo 數據集是一個常見用例，但模型預取不限於此數據集。

Fig. 7: Preprocessing of dataset for model oversubscription 圖來自源碼。

4.6 HugeCTR 到 ONNX 轉換器

HugeCTR to Open Neural Network Exchange (ONNX) 轉換器是一個hugectr2onnxPython 包，可以將 HugeCTR 模型轉換爲 ONNX。它可以提高 HugeCTR 與其他深度學習框架的兼容性，因爲 ONNX 作爲 AI 模型的開源格式。

使用我們的 HugeCTR Python API 進行訓練後，您可以獲得密集模型、稀疏模型和圖形配置的文件，這些文件在使用該hugectr2onnx.converter.convert方法時需要作爲輸入。每個 HugeCTR 層將對應一個或多個 ONNX 算子，訓練好的模型權重將作爲初始化器加載到 ONNX 圖中。此外，您可以選擇使用convert_embedding標誌轉換稀疏嵌入層。

4.7 分層參數服務器

HugeCTR 分層參數服務器 (POC) 上的本地 SSD 和 CPU 內存之間實現了分層存儲機制。通過這種實現，嵌入表不再需要存儲在本地 CPU 內存中。添加了分佈式 Redis 集羣作爲 CPU 緩存，以存儲更大的嵌入表並直接與 GPU 嵌入緩存交互。爲了幫助 Redis 集羣查找丟失的嵌入鍵，已實現本地 RocksDB 作爲查詢引擎來備份本地 SSD 上的完整嵌入表。

4.8 異步多線程數據管道

如果沒有高效的數據管道，即使向前和向後傳播以光速運行，其效果也如同到達機場的時間遠長於飛行時間。另外，當數據集很大並且經常變化時，將其拆分爲多個文件是非常合理的。

爲了有效地把數據獲取這個長延遲隱藏起來，HugeCTR 有一個多線程數據讀取器，其可以將數據獲取與實際模型訓練重疊起來。如下圖所示，DataReader是一個façade，由多個並行工作器和一個收集器組成。

每個工作器每次從其分配到的數據集文件中讀取一個批次。收集器會將收集到的數據記錄分發到多個 GPU。所有的工作人員、收集器和模型訓練作爲不同的線程在 CPU 上同時運行。

Figure 4. HugeCTR multithreaded data reader.

圖來自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

下圖顯示了 HugeCTR 流水線如何把 "數據從磁盤讀取到 CPU 內存的數據"，"從 CPU 到 GPU 的數據傳輸"以及"在 GPU 上跨不同批次的實際訓練"這三個階段重疊起來。

圖來自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

4.9 靈活模型配置

儘管 CTR 模型之間存在一些共性，但它們的細節（包括超參數）可能有所不同。爲了實現模型的靈活定製，HugeCTR 允許以 JSON 格式直觀地配置模型。

例如，要描述如下圖所示的混合模型，您可以編寫如圖 (b) 中抽象所示的“layers”子句。您可以有多個嵌入，您還可以指定批處理大小、優化器、數據路徑等。在同一個配置文件中，您也可以指定用於訓練的 GPU 數量和數量。有關更多信息，請參閱HugeCTR 用戶指南和示例配置文件。

Figure 6. A hybrid model with two embeddings and two different types of inputs. (a) An example mode expressible by HugeCTR. (b) The corresponding config. A lot of details are omitted for simplicity.

圖來自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems