[原始碼解析] 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 的神經網路進行互動。它還顯示瞭如何減少每個插槽的輸入特徵並跨兩個插槽連線

該圖顯示了一個跨越 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 上跨不同批次的實際訓練"這三個階段重疊起來。

此圖顯示了“讀取檔案”、“複製到 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