[原始碼解析] NVIDIA HugeCTR，GPU版本引數伺服器--- (5) 嵌入式hash表
0x00 摘要
在這篇文章中，我們介紹了 HugeCTR，這是一個面向行業的推薦系統訓練框架，針對具有模型並行嵌入和資料並行密集網路的大規模 CTR 模型進行了優化。
其中借鑑了HugeCTR原始碼閱讀 這篇大作，特此感謝。
本系列其他文章如下：
[原始碼解析] NVIDIA HugeCTR，GPU 版本引數伺服器 --(1)
[原始碼解析] NVIDIA HugeCTR，GPU版本引數伺服器--- (2)
[原始碼解析] NVIDIA HugeCTR，GPU版本引數伺服器---(3)
[原始碼解析] NVIDIA HugeCTR，GPU版本引數伺服器--- (4)
0x01 前文回顧
在前文，我們已經完成了對HugeCTR流水線的分析，接下來就要看嵌入層的實現，這部分是HugeCTR的精華所在。
嵌入在現代基於深度學習的推薦架構中發揮著關鍵作用，其為數十億實體（使用者、產品及其特徵）編碼個體資訊。隨著資料量的增加，嵌入表的大小也在增加，現在跨越多個 GB 到 TB。因為其巨大的嵌入表和稀疏訪問模式可能跨越多個 GPU，所以訓練這種型別的 DL 系統存在獨特的挑戰。
HugeCTR 就實現了一種優化的嵌入實現，其效能比其他框架的嵌入層高 8 倍。這種優化的實現也可作為 TensorFlow 外掛使用，可與 TensorFlow 無縫協作，並作為 TensorFlow 原生嵌入層的便捷替代品。
0x02 Embedding
2.1 概念
我們先簡要介紹一下embedding的概念。嵌入（embedding）是一種機器學習技術，被用來把一個id/category特徵自動轉換為待優化的特徵向量，這樣可以把演算法從精確匹配擴充到模糊匹配，從而提高演算法的擴充能力。從另一個角度看，嵌入表是一種特定型別的key-value儲存，鍵是用於唯一標識物件的 ID，值是實數向量。
Embedding技術在 NLP 中很流行，用來把單詞表示密集數值向量，具有相似含義的單詞具有相似的嵌入向量。
另外，Embedding 還有一個優勢就是把資料從高維轉換為低維。
2.1.1 One-hot 編碼
作為比對，我們先看看 One-hot 編碼。One-hot編碼就是保證每個樣本中的單個特徵只有1位處於狀態1，其他的都是0。具體編碼舉例如下，把語料庫中，杭州、上海、寧波、北京每個都對應一個向量，向量中只有一個值為1，其餘都為0。
杭州 [0,0,0,0,0,0,0,1,0,……，0,0,0,0,0,0,0]
上海 [0,0,0,0,1,0,0,0,0,……，0,0,0,0,0,0,0]
寧波 [0,0,0,1,0,0,0,0,0,……，0,0,0,0,0,0,0]
北京 [0,0,0,0,0,0,0,0,0,……，1,0,0,0,0,0,0]
其缺點是：
- 向量的維度會隨著詞的數量增大而增大；如果將世界所有城市名稱對應的向量組成一個矩陣的話，那這個矩陣因為過於稀疏則會造成維度災難。
- 因為城市編碼是隨機的，所以向量之間相互獨立，無法表示詞彙之間在語義層面上的相關資訊。
所以，人們想對獨熱編碼做如下改進：
- 將vector每一個元素由整型改為浮點型，從整型變為整個實數範圍的表示；
- 把原始稀疏向量轉化為低維度的連續值，也就是稠密向量。可以認為是將原來稀疏的巨大維度壓縮嵌入到一個更小維度的空間。並且其中意思相近的詞將被對映到這個向量空間中相近的位置。
簡單說，就是尋找一個空間對映把高維詞向量嵌入到一個低維空間。
2.1.2 分散式表示
分散式表示（Distributed Representation）基本思想是將每個詞表達成 n 維稠密、連續的實數向量。而實數向量之間的關係可以代表詞語之間的相似度，比如向量的夾角cosine或者歐氏距離。用詞彙舉例，獨熱編碼相當於對詞進行編碼，而分散式表示則是將詞從稀疏的大維度壓縮嵌入到較低維度的向量空間中。
分散式表示最大的貢獻就是讓相關或者相似的詞在距離上更接近了。分散式表示相較於One-hot方式另一個區別是維數下降很多，對於一個100萬的詞表，我們可以用100維的實數向量來表示一個詞，而One-hot得要100W萬來編碼。比如杭州、上海、寧波、北京，廣州，深圳，瀋陽，西安，洛陽 這九個城市 採用 one-hot 編碼如下：
杭州 [1,0,0,0,0,0,0,1,0]
上海 [0,1,0,0,0,0,0,0,0]
寧波 [0,0,1,0,0,0,0,0,0]
北京 [0,0,0,1,0,0,0,0,0]
廣州 [0,0,0,0,1,0,0,0,0]
深圳 [0,0,0,0,0,1,0,0,0]
瀋陽 [0,0,0,0,0,0,1,0,0]
西安 [0,0,0,0,0,0,0,1,0]
洛陽 [0,0,0,0,0,0,0,0,1]
但是太稀疏，佔用太大記憶體，所以弄一個稠密矩陣，也就是 Embedding Table 如下：
如果要查詢杭州，上海：
利用矩陣乘法，可以得到：
這樣就把兩個 1 x 9 的高維度，離散，稀疏向量，壓縮到 兩個 1 x 3 的低維稠密向量。這裡把 One-Hot 向量中 “1”的位置叫做sparseID，就是一個編號。這個獨熱向量和嵌入表的矩陣乘法就等於利用sparseID進行的一次查表過程，就是依據杭州，上海的sparseID（0，1），從嵌入表中取出對應的向量（第0行、第1行）。這樣就把高維變成了低維。
2.1.3 推薦領域
在推薦系統領域，Embedding將每個感興趣的物件（使用者、產品、類別等）表示為一個密集的數值向量。最簡單的推薦系統基於使用者和產品：您應該向使用者推薦哪些產品？您有使用者 ID 和產品 ID作為key。對應的value則是使用者和產品，因此您使用兩個嵌入表。
圖1. 嵌入表是稀疏類別的密集表示。 每個類別由一個向量表示，這裡嵌入維度是4。來自 Using Neural Networks for Your Recommender System。
2.2 Lookup
從上圖能夠看到另外一個概念：lookup，我們接下來就分析一下。
Embeddings層的神經元個數是由embedding vector和field_size共同確定，即神經元的個數為embedding vector * field_size。dense vector是embeding層的輸出的水平拼接。由於輸入特徵one-hot編碼，所以embedding vector也就是輸入層到Dense Embeddings層的權重，也就是全連線層的權重。
Embedding 權重矩陣可以是一個 [item_size, embedding_size] 的稠密矩陣，item_size是需要embedding的物品個數，embedding_size是對映的向量長度，或者說矩陣的大小是：特徵數量 * 嵌入維度。Embedding 權重矩陣的每一行對應輸入的一個維度特徵（one-hot之後的維度）。使用者可以用一個index表示選擇了哪個特徵。
Embedding_lookup 就是如何從Embedding 權重矩陣獲取到一個超高維輸入對應的embedding向量的方法，embedding就是權重本身。Embedding_lookup 實際上是由矩陣相乘實現的 V = WX + b，因為輸入 X 是One-Hot編碼，所以和矩陣相乘相當於是取出權重矩陣中對應的那一行，看起來像是在查一個索引表，所以叫做 lookup。Embedding 本質上可以看做是一個全連線層。比如：
嵌入層這些權重是通過神經網路自己學習到的，實際上，權重矩陣一般是隨機初始化的，是需要優化的變數。訓練神經網路時，每個Embedding向量都會得到更新，即在不斷升維和降維的過程中，找到最適合的維度。因此，embedding_lookup 這裡還需要完成反向傳播，即自動求導和對權重矩陣的更新。
2.3 嵌入層
嵌入層是現代深度學習推薦系統的關鍵模組，其通常位於輸入層之後，在特徵互動和密集層之前。嵌入層就像深度神經網路的其他層一樣，是從資料和端到端訓練中學習得到的。我們接下來看看如何使用嵌入層。
2.3.1 點積
我們可以計算使用者嵌入和專案嵌入之間的點積以獲得最終分數，即使用者與專案互動的可能性。您可以應用 sigmoid 啟用函式作為將輸出轉換為 0 到 1 之間的概率的最後一步。
圖 2. 具有兩個嵌入表和點積輸出的神經網路。來自 Using Neural Networks for Your Recommender System。
此方法等效於矩陣分解或交替最小二乘法 (ALS)。
2.3.2 全連線層
如果使用多個非線性層構建一個深層結構，則神經網路效能會更佳。您可以通過使用 ReLU 啟用將嵌入層的輸出饋送到多個完全連線的層來擴充套件先前的模型。這裡在設計上的一個選擇點是：如何組合兩個嵌入向量。您可以只連線嵌入向量，也可以將向量按元素相乘，類似於點積。點積輸出後跟著多個隱藏層。
圖 3. 具有兩個嵌入表和多個全連線層的神經網路，將使用者和產品嵌入進行concatenate 或者進行元素級別（element-wise）相乘。 多個隱藏層可以處理結果向量。來自 Using Neural Networks for Your Recommender System。
2.3.3 後設資料資訊
到目前為止，我們只使用了使用者 ID 和產品 ID 作為輸入，但我們通常有更多可用資訊。比如使用者的其他資訊可以是性別、年齡、城市（地址）、自上次訪問以來的時間或用於付款的信用卡。一件商品通常有過去 7 天內售出的品牌、價格、類別或數量。這些輔助資訊可以幫助模型更好地泛化。我們可以修改神經網路以使用附加特徵作為輸入。
圖 4. 具有元資訊和多個全連線層的神經網路，我們向神經網路架構新增更多資訊，比如可以新增例如城市、年齡、分支機構、類別和價格等輔助資訊。 來自 Using Neural Networks for Your Recommender System。
至此，我們可知，基於嵌入層我們可以得到一個基礎的推薦系統網路。
2.3.4 經典架構
接下來，我們看看 2016 年 Google 的 Wide and Deep 和 Facebook 2019 年的 DLRM。
2.3.4.1 Google’s Wide and Deep
Google 的 Wide 和 Deep 包含兩個元件：
- 與之前的神經網路相比，Wide部分是新元件，它是輸入特徵的線性組合，具有類似線性/邏輯迴歸的邏輯。
- Deep部分的作用是：把高維，稀疏的類別特徵通過嵌入層處理為低維稠密向量，並將輸出與連續值特徵連線起來。連線之後的向量傳遞到MLP層。
兩個部分的輸出通過加權相加合併在一起得到最終預測值。
對於推薦系統模型來說，理想狀態是同時具有記憶能力和泛化能力。
-
泛化能力指的是模型對未觀察到的特徵或者罕見特徵可以做出預測。
- 比如通過學習得到一個規則是：有翅膀會飛，所以看見喜鵲就可以推斷它會飛，這就是泛化能力。
- Deep部分 通過基於使用者/物品特徵進行歸納以提供優秀的泛化能力。
-
記憶能力指的是模型可以記住大量的歷史行為特徵，然後從歷史資料之中學習到特徵的共性，而且可以把這些共性作為推薦依據。
- 泛化能力會有疏漏，比如企鵝不會飛。
- Wide部分 就可以通過歷史等記憶資訊來提供優秀的記憶能力，糾正例外情況。邏輯迴歸這樣的簡單模型如果發現一個“強特徵”，就會在訓練時候把其相應的權重調整得非常大，從而實現了對這個特徵的記憶，這就是所謂的模型“記憶能力”。
- 我們給出論文（https://arxiv.org/abs/1606.07792）之中的架構圖，可以看到，wide部分選擇了兩個特徵：使用者安裝的app，曝光app，然後加交叉得到了一個特徵，比如使用者已經安裝了netflix，而且在應用商店中曾經看過pandora應用，那麼交叉特徵：AND(user_installed_app=netflix, impression_app=pandora)就會被設定為1。這就是希望模型可以記住這樣的規則：“如果使用者已經安裝了應用 A，是否會安裝 B”。
因此，Wide and Deep 模型兼具邏輯迴歸和深度神經網路的優點，可以記憶大量歷史行為，又擁有強大的表達能力。
在HugeCTR之中，可以看到是通過如下方式來組織模型的。
2.3.4.2 Facebook 的 DLRM
Facebook 的 DLRM（Deep Learning Recommendation Model） 與帶有後設資料的神經網路架構具有相似的結構，但有一些區別。資料集可以包含多個分類特徵。DLRM 要求所有分類輸入都通過具有相同維度的嵌入層饋送。
接下來，連續輸入被串聯並通過多個完全連線的層饋送，稱為底層多層感知器 (MLP)。底部 MLP 的最後一層與嵌入層向量具有相同的維度。
DLRM 使用新的組合層。它在所有嵌入向量對和底部 MLP 輸出之間應用逐元素（element-wise）乘法。這就是每個向量具有相同維度的原因。生成的向量被連線起來並且傳送給另一組全連線層（頂部 MLP）。
圖 5. Wide and Deep 架構在左側視覺化，DLRM 架構在右側。
本節圖片來自 Using Neural Networks for Your Recommender System。
2.4 推薦系統的嵌入層
因為CTR領域之中，特徵的特點是高維，稀疏，所以對於離散特徵，一般使用one-hot編碼。但是將One-hot型別的特徵輸入到DNN中，會導致網路引數太多，比如輸入層有1000萬個節點，隱層有500節點，則引數有50億個。
所以人們增加了一個Embedding層用於降低維度，這樣就對單個特徵的稀疏向量進行緊湊化處理。但是有時還是不奏效，人們也可以將特徵分為不同的field。
2.4.1 特色
和與其他型別 DL 模型相比，DL 推薦模型的嵌入層是比較特殊的：它們為模型貢獻了大量引數，但幾乎不需要計算，而計算密集型denser layers的引數數量則要少得多。
舉一個具體的例子：原始的Wide and Deep模型有幾個大小為[1024,512,256]的dense layers，因此這些dense layers只有幾百萬個引數，而其嵌入層可以有數十億個條目，以及數十億個引數。這與例如在 NLP 領域流行的 BERT 模型架構形成對比，BERT的嵌入層只有數萬個條目，總計數百萬個引數，但其稠密的前饋和注意力層則由數億個引數組成。這種差異還導致另一個結果：與其他型別的 DL 模型相比，DL 推薦網路輸入資料的每位元組計算量通常要小得多。
2.4.2 優化嵌入重要性
對於推薦系統，嵌入層的優化十分重要。要理解為什麼嵌入層和相關操作的優化很重要，首先要看看推薦系統嵌入層訓練所遇到的挑戰：資料量和速度。
2.4.2.1 資料量
隨著線上平臺和服務獲得數億甚至數十億使用者，並且提供的獨特產品數量達到數十億，嵌入表的規模越來越大也就不足為奇了。
據報導，Instagram 一直致力於開發大小達到 10 TB 的推薦模型。同樣，百度的一個廣告排名模型，也達到了 10 TB 的境界。在整個行業中，數百 GB 到 TB 的模型正變得越來越流行，例如Pinterest 的 4-TB 模型和Google 的 1.2-TB 模型。因此，在單個計算節點上擬合 TB 級模型是一項重大挑戰，更不用說在單個計算加速器（比如GPU）了。NVIDIA A100 GPU 只是配備了 80 GB 的 HBM。
2.4.2.2 訪問速度
訓練推薦系統本質上是一項記憶體頻寬密集型任務。這是因為每個訓練樣本或批次通常涉及嵌入表中的少量實體。必須檢索這些條目才能計算前向傳遞，然後在後向傳遞中更新。
CPU 主存容量大但頻寬有限，高階機型通常在幾十 GB/s 範圍內。另一方面，GPU 的記憶體容量有限，但頻寬很高。NVIDIA A100 80-GB GPU 提供 2 TB/s 的記憶體頻寬。
2.4.3 解決方案
針對這些挑戰，人們已經找到了一些解決方案，但是都存在一些問題，比如：
- 將整個嵌入表儲存在主儲存器上解決了大小問題。然而，它通常會導致訓練吞吐量極慢，而新資料的數量和速度往往使這種吞吐量相形見絀，從而導致系統無法及時重新訓練。
- 或者，可以把嵌入層分佈在多個 GPU 和多個節點上，但這樣卻陷入通訊瓶頸，導致 GPU 計算利用率不足，訓練效能與純 CPU 訓練不相上下。
因此，嵌入層是推薦系統的主要瓶頸之一。優化嵌入層是解鎖 GPU 高計算吞吐量的關鍵。
0x03 DeepFM
因為我們要用DeepFM作為示例，所以需要介紹一下基本內容。
我們選擇 IJCAI 2017 的論文 DeepFM: A Factorization-Machine based Neural Network for CTR Prediction 的內容做分析。
3.1 CTR特點
CTR預估資料有如下特點：
- 輸入的資料有類別型和連續型。類別型資料會編碼成one-hot，連續型資料可以先離散化再變嗎為one-hot，也可以保留原值。
- 資料的維度非常高。
- 資料非常稀疏。
- 特徵按照Field分組。
CTR預估重點在於學習組合特徵。Google論文研究結論為：高階和低階的組合特徵都非常重要，應該同時學習到這兩種組合特徵。所以關鍵點是如何高效的提取這些組合特徵。
3.2 DeepFM
DeepFM將Google的wide & deep模型進行改進：
- 將Wide & Deep 部分的wide部分由 人工特徵工程 + LR 轉換為FM模型，FM提取低階組合特徵，Deep提取高階組合特徵，這樣避開了人工特徵工程，提高了模型的泛化能力；
- FM模型和Deep部分共享Embedding，使得DeepFM成為一種端到端的模型，提高了模型的訓練效率，不但訓練更快而且更準確。
具體模型架構如下：
0x04 HugeCTR嵌入層
為了克服嵌入挑戰並實現更快的訓練，HugeCTR 實現了自己的嵌入層，其中包括一個 GPU 加速的雜湊表、以節省記憶體的方式實現的高效稀疏優化器以及各種嵌入分佈策略。它利用NCCL作為其 GPU 間通訊原語。
4.1 雜湊表
雜湊表的實現基於RAPIDS cuDF，它是一個 GPU DataFrame 庫，是 NVIDIA 的 RAPIDS 資料科學平臺的一部分。cuDF GPU 雜湊表可以比基於 Threading Building Blocks (TBB) 實現的 concurrent_hash_map 更快，而且達到 35 倍的加速。
4.2 模型並行
HugeCTR 提供了一個模型並行的嵌入表，其分佈在叢集中的所有 GPU 上，叢集由多個節點和多個 GPU 組成。另一方面，密集層採用資料並行性，每個 GPU 上有一個副本。
考慮到可擴充套件性，HugeCTR 預設支援嵌入層的模型並行性。
圖 6. HugeCTR 模型和資料並行架構。來自 https://developer.nvidia.com/blog/announcing-nvidia-merlin-application-framework-for-deep-recommender-systems/。
4.3 通訊
HugeCTR 使用NCCL 完成了高速和可擴充套件的節點間和節點內通訊。對於有很多輸入特徵的情況，HugeCTR嵌入表可以分割成多個槽。屬於同一個槽的特徵被獨立轉換為對應的嵌入向量，然後被規約為單個嵌入向量。這允許使用者將每個插槽中的有效功能的數量有效地減少到可管理的程度。
4.4 印證
我們可以從 https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/Recommendation/DLRM#hybrid-parallel-multi-gpu-with-all-2-all-communication 之中看到混合並行的思路，可以和HugeCTR印證（因為我們使用DeepFM來剖析，所以DLRM只用於印證）。
4.4.1 DLRM
在DLRM之中，為了處理類別資料，嵌入層將每個類別對映到密集表示，然後再輸入多層感知器 (MLP)。數值特徵則可以直接輸入 MLP。在下一級，通過取所有嵌入向量對和處理過的密集特徵之間的點積，顯式計算不同特徵的二階互動。這些成對的互動被饋送到頂級 MLP 以計算使用者和產品對之間互動的可能性。
與其他基於 DL 的推薦方法相比，DLRM 在兩個方面有所不同。首先，它明確計算了特徵互動，同時將互動順序限制為成對互動（pairwise interactions）。其次，DLRM 將每個嵌入的特徵向量（對應於類別特徵）視為一個單元，而其他方法（例如 Deep and Cross）將特徵向量中的每個元素視為一個新單元，這樣會會產生不同的交叉項。這些設計選擇有助於降低計算和記憶體成本，同時也可以提供相當的準確性。
4.4.2 混合並行
許多推薦模型包含非常大的嵌入表。因此，該模型往往太大，無法裝入單個裝置。這可以通過使用CPU或其他GPU作為 "記憶體捐贈者(memory donors)"，以模型並行的方式進行訓練而輕鬆解決。然而，這種方法是次優的，因為 "記憶體捐贈者 "裝置的計算沒有被利用。
針對 DLRM，我們對模型的底層部分使用模型並行方法（嵌入表+底層MLP），而對模型的頂層部分使用通常的資料並行方法（Dot Interaction + Top MLP）。這樣，我們可以訓練比通常適合單個GPU的模型大得多的模型，同時通過使用多個GPU使訓練更快。我們稱這種方法為混合並行。
4.5 使用
我們可以採用如下兩種方式來使用 hugeCTR 嵌入層：
- 將原生 NVIDIA Merlin HugeCTR 框架用於訓練和推理工作。
- 使用 NVIDIA Merlin HugeCTR TensorFlow 外掛，該外掛旨在與 TensorFlow 無縫協作。
0x05 初始化
5.1 配置
前面提到了可以使用程式碼完成網路配置，我們從下面可以看到，DeepFM 一共有三個embedding層，分別對應 wide_data 的 sparse 引數對映到dense vector，deep_data 的 sparse 引數對映到dense vector，。
model = hugectr.Model(solver, reader, optimizer)
model.add(hugectr.Input(label_dim = 1, label_name = "label",
dense_dim = 13, dense_name = "dense",
data_reader_sparse_param_array =
[hugectr.DataReaderSparseParam("wide_data", 30, True, 1),
hugectr.DataReaderSparseParam("deep_data", 2, False, 26)]))
model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash,
workspace_size_per_gpu_in_mb = 23,
embedding_vec_size = 1,
combiner = "sum",
sparse_embedding_name = "sparse_embedding2",
bottom_name = "wide_data",
optimizer = optimizer))
model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash,
workspace_size_per_gpu_in_mb = 358,
embedding_vec_size = 16,
combiner = "sum",
sparse_embedding_name = "sparse_embedding1",
bottom_name = "deep_data",
optimizer = optimizer))
5.1.1 DataReaderSparseParam
DataReaderSparseParam 定義如下，其中slot_num就代表了這個層擁有幾個slot。比如 hugectr.DataReaderSparseParam("deep_data", 2, False, 26) 就代表了有26個slots。
struct DataReaderSparseParam {
std::string top_name;
std::vector<int> nnz_per_slot;
bool is_fixed_length;
int slot_num;
DataReaderSparse_t type;
int max_feature_num;
int max_nnz;
DataReaderSparseParam() {}
DataReaderSparseParam(const std::string& top_name_, const std::vector<int>& nnz_per_slot_,
bool is_fixed_length_, int slot_num_)
: top_name(top_name_),
nnz_per_slot(nnz_per_slot_),
is_fixed_length(is_fixed_length_),
slot_num(slot_num_),
type(DataReaderSparse_t::Distributed) {
if (static_cast<size_t>(slot_num_) != nnz_per_slot_.size()) {
CK_THROW_(Error_t::WrongInput, "slot num != nnz_per_slot.size().");
}
max_feature_num = std::accumulate(nnz_per_slot.begin(), nnz_per_slot.end(), 0);
max_nnz = *std::max_element(nnz_per_slot.begin(), nnz_per_slot.end());
}
DataReaderSparseParam(const std::string& top_name_, const int nnz_per_slot_,
bool is_fixed_length_, int slot_num_)
: top_name(top_name_),
nnz_per_slot(slot_num_, nnz_per_slot_),
is_fixed_length(is_fixed_length_),
slot_num(slot_num_),
type(DataReaderSparse_t::Distributed) {
max_feature_num = std::accumulate(nnz_per_slot.begin(), nnz_per_slot.end(), 0);
max_nnz = *std::max_element(nnz_per_slot.begin(), nnz_per_slot.end());
}
};
5.1.2 slot概念
我們從文件之中可以知道，slot的概念就是特徵域或者表。
In HugeCTR, a slot is a feature field or table. The features in a slot can be one-hot or multi-hot. The number of features in different slots can be various. You can specify the number of slots (`slot_num`) in the data layer of your configuration file.
Field 或者 slots（有的文章也稱其為Feature Group）就是若干有關聯特徵的集合，其主要作用就是把相關特徵組成一個feature field，然後把這個field再轉換為一個稠密向量，這樣就可以減少DNN輸入層規模和模型引數。
比如：使用者看過的商品，使用者購買的商品，這就是兩個Field，具體每個商品則是feature，這些商品共享一個商品列表，或者說共享一份Vocabulary。如果把每個商品都進行embedding，然後把這些張量拼接起來，那麼DNN輸入層就太大了。所以把這些購買的商品歸類為同一個field，把這些商品的embedding向量再做pooling之後得到一個field的向量，那麼輸入層數目就少了很多。
5.1.2.1 FLEN
FLEN: Leveraging Field for Scalable CTR Prediction 之中有一些論證和幾張精彩圖例可以來輔助說明這個概念，具體如下：
CTR預測任務中的資料是多域（multi-field categorical ）的分類資料，也就是說，每個特徵都是分類的，並且只屬於一個欄位。例如，特徵 "gender=Female "屬於域 "gender"，特徵 "age=24 "屬於域 "age"，特徵 "item category=cosmetics "屬於域 "item category"。特徵 "性別 "的值是 "男性 "或 "女性"。特徵 "年齡 "被劃分為幾個年齡組。"0-18歲"，"18-25歲"，"25-30歲"，等等。人們普遍認為，特徵連線（conjunctions）是準確預測點選率的關鍵。一個有資訊量的特徵連線的例子是：年齡組 "18-25 "與性別 "女性 "相結合，用於 "化妝品 "專案類別。它表明，年輕女孩更有可能點選化妝品產品。
FLEN模型中使用了一個filed-wise embedding vector，通過將同一個域（如user filed或者item field）之中的embedding 向量進行求和，來得到域對應的embedding向量。
比如，首先把特徵\(x_n\) 轉換為嵌入向量 \(e_n\)。
其次，使用 sum-pooling 得到 field-wise embedding vectors。
比如
最後，把所有field-wise embedding vectors拼接起來。
系統整體架構如下：
5.1.2.2 Pooling
具體如何做pooling？HugeCTR有sum或者mean兩種操作，具體叫做combiner，比如：
// do sum reduction
if (combiner == 0) {
forward_sum(batch_size, slot_num, embedding_vec_size, row_offset.get_ptr(),
hash_value_index.get_ptr(), hash_table_value.get_ptr(),
embedding_feature.get_ptr(), stream);
} else if (combiner == 1) {
forward_mean(batch_size, slot_num, embedding_vec_size, row_offset.get_ptr(),
hash_value_index.get_ptr(), hash_table_value.get_ptr(),
embedding_feature.get_ptr(), stream);
}
結合前面圖，類別特徵就一共有M個slot，對應了M個嵌入表。
比如在 test/pybind_test/din_matmul_fp32_1gpu.py 之中，可見CateID有11個slots。
model.add(hugectr.Input(label_dim = 1, label_name = "label",
dense_dim = 0, dense_name = "dense",
data_reader_sparse_param_array =
[hugectr.DataReaderSparseParam("UserID", 1, True, 1),
hugectr.DataReaderSparseParam("GoodID", 1, True, 11),
hugectr.DataReaderSparseParam("CateID", 1, True, 11)]))
比如，下圖之中，一個sample有7個key，分為兩個field，就是兩個slot。4個key放在第一個slot之上，3個key放到第二個slot上，第三個slot沒有key。在查詢過程中，會把這些key對應的value查詢出來。第一個slot內部會對這些value進行sum或者mean操作得到V1，第二個slot內部對3個value進行sum或者mean操作得到V2，最後會把V1，V2進行concat操作，傳送給後續層。
5.1.2.3 TensorFlow
我們可以從TF原始碼註釋裡面看到pooling的一些使用。
tensorflow/python/ops/embedding_ops.py 是關於embedding的使用。
combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
and "sum" are supported. "sum" computes the weighted sum of the embedding
results for each row. "mean" is the weighted sum divided by the total
weight. "sqrtn" is the weighted sum divided by the square root of the sum
of the squares of the weights. Defaults to `mean`.
tensorflow/python/feature_column/feature_column.py 是關於feature column的使用。
sparse_combiner: A string specifying how to reduce if a categorical column
is multivalent. Except `numeric_column`, almost all columns passed to
`linear_model` are considered as categorical columns. It combines each
categorical column independently. Currently "mean", "sqrtn" and "sum" are
supported, with "sum" the default for linear model. "sqrtn" often achieves
good accuracy, in particular with bag-of-words columns.
* "sum": do not normalize features in the column
* "mean": do l1 normalization on features in the column
* "sqrtn": do l2 normalization on features in the column
在：tensorflow/lite/kernels/embedding_lookup_sparse.cc 的註釋有直接從嵌入表之中look up的使用。
// Op that looks up items from a sparse tensor in an embedding matrix.
// The sparse lookup tensor is represented by three individual tensors: lookup,
// indices, and dense_shape. The representation assume that the corresponding
// dense tensor would satisfy:
// * dense.shape = dense_shape
// * dense[tuple(indices[i])] = lookup[i]
//
// By convention, indices should be sorted.
//
// Options:
// combiner: The reduction op (SUM, MEAN, SQRTN).
// * SUM computes the weighted sum of the embedding results.
// * MEAN is the weighted sum divided by the total weight.
// * SQRTN is the weighted sum divided by the square root of the sum of the
// squares of the weights.
//
// Input:
// Tensor[0]: Ids to lookup, dim.size == 1, int32.
// Tensor[1]: Indices, int32.
// Tensor[2]: Dense shape, int32.
// Tensor[3]: Weights to use for aggregation, float.
// Tensor[4]: Params, a matrix of multi-dimensional items,
// dim.size >= 2, float.
//
// Output:
// A (dense) tensor representing the combined embeddings for the sparse ids.
// For each row in the sparse tensor represented by (lookup, indices, shape)
// the op looks up the embeddings for all ids in that row, multiplies them by
// the corresponding weight, and combines these embeddings as specified in the
// last dimension.
//
// Output.dim = [l0, ... , ln-1, e1, ..., em]
// Where dense_shape == [l0, ..., ln] and Tensor[4].dim == [e0, e1, ..., em]
//
// For instance, if params is a 10x20 matrix and ids, weights are:
//
// [0, 0]: id 1, weight 2.0
// [0, 1]: id 3, weight 0.5
// [1, 0]: id 0, weight 1.0
// [2, 3]: id 1, weight 3.0
//
// with combiner=MEAN, then the output will be a (3, 20) tensor where:
//
// output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
// output[1, :] = (params[0, :] * 1.0) / 1.0
// output[2, :] = (params[1, :] * 3.0) / 3.0
//
// When indices are out of bound, the op will not succeed.
另外，其他框架/模型實現也有使用加權平均（比如使用Attention），或者加入時序資訊。
5.2 構建
這是一個比較複雜的過程，從前文我們知道，DataReader最後把各種輸入都拷貝到了其成員變數 output_ 之上。
那麼，嵌入層是如何利用到 output_ 中的sparse特徵的呢？我們需要一步一步來看。
5.2.1 流水線
parser.cpp 之中，如下程式碼建立了流水線，我們省略了很多程式碼，可以看到先呼叫 create_datareader 建立了reader，然後才建立 embedding。
template <typename TypeKey>
void Parser::create_pipeline_internal(std::shared_ptr<IDataReader>& init_data_reader,
std::shared_ptr<IDataReader>& train_data_reader,
std::shared_ptr<IDataReader>& evaluate_data_reader,
std::vector<std::shared_ptr<IEmbedding>>& embeddings,
std::vector<std::shared_ptr<Network>>& networks,
const std::shared_ptr<ResourceManager>& resource_manager,
std::shared_ptr<ExchangeWgrad>& exchange_wgrad) {
try {
std::map<std::string, SparseInput<TypeKey>> sparse_input_map;
{
// Create Data Reader
{
create_datareader<TypeKey>()(j, sparse_input_map, train_tensor_entries_list,
evaluate_tensor_entries_list, init_data_reader,
train_data_reader, evaluate_data_reader, batch_size_,
batch_size_eval_, use_mixed_precision_, repeat_dataset_,
enable_overlap, resource_manager);
} // Create Data Reader
// Create Embedding ---- 這裡建立了embedding
{
for (unsigned int i = 1; i < j_layers_array.size(); i++) {
if (use_mixed_precision_) {
create_embedding<TypeKey, __half>()(
sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list,
embeddings, embedding_type, config_, resource_manager, batch_size_,
batch_size_eval_, exchange_wgrad, use_mixed_precision_, scaler_, j, use_cuda_graph_,
grouped_all_reduce_);
} else {
create_embedding<TypeKey, float>()(
sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list,
embeddings, embedding_type, config_, resource_manager, batch_size_,
batch_size_eval_, exchange_wgrad, use_mixed_precision_, scaler_, j, use_cuda_graph_,
grouped_all_reduce_);
}
} // for ()
} // Create Embedding
// create network
for (size_t i = 0; i < resource_manager->get_local_gpu_count(); i++) {
networks.emplace_back(Network::create_network(
j_layers_array, j_optimizer, train_tensor_entries_list[i],
evaluate_tensor_entries_list[i], total_gpu_count, exchange_wgrad,
resource_manager->get_local_cpu(), resource_manager->get_local_gpu(i),
use_mixed_precision_, enable_tf32_compute_, scaler_, use_algorithm_search_,
use_cuda_graph_, false, grouped_all_reduce_));
}
}
exchange_wgrad->allocate();
} catch (const std::runtime_error& rt_err) {
std::cerr << rt_err.what() << std::endl;
throw;
}
}
5.2.2 建立 DataReader
我們來看看如何建立 Reader，依然省略大部分程式碼。主要和sparse輸入相關的有如下：
- create_datareader 的引數有sparse_input_map，這是一個引用。
- 依據配置對 sparse_input_map 進行設定。
- 依據引數名字找到 sparse_input。
- 把 reader 的output_ 之中的 sparse_tensors_map 賦值到 sparse_input 之中。就是這行程式碼
copy(data_reader_tk->get_sparse_tensors(sparse_name), sparse_input->second.train_sparse_tensors);把 reader 的 output_ 之中的 sparse_tensors_map 賦值到 sparse_input 之中。 - 所以reader裡面最終設定的是 sparse_input_map 之中某一個 sparse_input->second.train_sparse_tensors。
- sparse_input_map 接下來就被傳入到 create_embedding 之中。
其中關鍵所在見下面程式碼註釋，
template <typename TypeKey>
void create_datareader<TypeKey>::operator()(
const nlohmann::json& j, std::map<std::string, SparseInput<TypeKey>>& sparse_input_map,
std::vector<TensorEntry>* train_tensor_entries_list,
std::vector<TensorEntry>* evaluate_tensor_entries_list,
std::shared_ptr<IDataReader>& init_data_reader, std::shared_ptr<IDataReader>& train_data_reader,
std::shared_ptr<IDataReader>& evaluate_data_reader, size_t batch_size, size_t batch_size_eval,
bool use_mixed_precision, bool repeat_dataset, bool enable_overlap,
const std::shared_ptr<ResourceManager> resource_manager) {
std::vector<DataReaderSparseParam> data_reader_sparse_param_array;
// 依據配置對 sparse_input_map 進行設定
for (unsigned int i = 0; i < j_sparse.size(); i++) {
DataReaderSparseParam param{sparse_name, nnz_per_slot_vec, is_fixed_length, slot_num};
data_reader_sparse_param_array.push_back(param);
SparseInput<TypeKey> sparse_input(param.slot_num, param.max_feature_num);
sparse_input_map.emplace(sparse_name, sparse_input);
sparse_names.push_back(sparse_name);
}
if (format == DataReaderType_t::RawAsync) {
// 省略
} else {
DataReader<TypeKey>* data_reader_tk = new DataReader<TypeKey>(......);
for (unsigned int i = 0; i < j_sparse.size(); i++) {
const std::string& sparse_name = sparse_names[i];
// 根據名字找到sparse輸入
const auto& sparse_input = sparse_input_map.find(sparse_name);
auto copy = [](const std::vector<SparseTensorBag>& tensorbags,
SparseTensors<TypeKey>& sparse_tensors) {
sparse_tensors.resize(tensorbags.size());
for (size_t j = 0; j < tensorbags.size(); ++j) {
sparse_tensors[j] = SparseTensor<TypeKey>::stretch_from(tensorbags[j]);
}
};
// 關鍵所在，把 reader 的output_ 之中的 sparse_tensors_map 賦值到 sparse_input 之中
copy(data_reader_tk->get_sparse_tensors(sparse_name),
sparse_input->second.train_sparse_tensors);
copy(data_reader_eval_tk->get_sparse_tensors(sparse_name),
sparse_input->second.evaluate_sparse_tensors);
}
}
get_sparse_tensors 程式碼如下：
const std::vector<SparseTensorBag> &get_sparse_tensors(const std::string &name) {
if (output_->sparse_tensors_map.find(name) == output_->sparse_tensors_map.end()) {
CK_THROW_(Error_t::IllegalCall, "no such sparse output in data reader:" + name);
}
return output_->sparse_tensors_map[name];
}
所以邏輯擴充如下：
- a) create_pipeline_internal 生成了 sparse_input_map，作為引數傳遞給 create_datareader；
- b) create_datareader 之中對sparse_input_map進行操作，使其指向了 DataReader.output_.sparse_tensors_map；
- c) sparse_input_map 作為引數傳遞給 create_embedding，具體如下：
create_embedding<TypeKey, float>()(sparse_input_map,......) - d) 所以，create_embedding 最終用到的就是 GPU 之上的sparse_input_map;
5.2.3 建立嵌入
如下程式碼建立了嵌入。
create_embedding<TypeKey, float>()(
sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list,
embeddings, embedding_type, config_, resource_manager, batch_size_,
batch_size_eval_, exchange_wgrad, use_mixed_precision_, scaler_, j, use_cuda_graph_,
grouped_all_reduce_);
這裡建立了一些embedding，比如DistributedSlotSparseEmbeddingHash，這些embedding最後儲存在 Session 的 embeddings_ 成員變數之中。這裡關鍵有幾點：
- 從sparse_input_map之中獲取sparse_input資訊。
- 使用sparse.train_sparse_tensors來構建 DistributedSlotSparseEmbeddingHash。所以我們可以知道，DataReader.output_ 成員變數將要和 DistributedSlotSparseEmbeddingHash 聯絡到一起。這裡是embedding的輸入。
- 輸出引數 train_tensor_entries_list 會作為 embedding 的輸出返回，這是一個指標。
template <typename TypeKey, typename TypeFP>
static void create_embeddings(std::map<std::string, SparseInput<TypeKey>>& sparse_input_map,
std::vector<TensorEntry>* train_tensor_entries_list,
std::vector<TensorEntry>* evaluate_tensor_entries_list,
std::vector<std::shared_ptr<IEmbedding>>& embeddings,
Embedding_t embedding_type, const nlohmann::json& config,
const std::shared_ptr<ResourceManager>& resource_manager,
size_t batch_size, size_t batch_size_eval,
std::shared_ptr<ExchangeWgrad>& exchange_wgrad,
bool use_mixed_precision,
float scaler, const nlohmann::json& j_layers,
bool use_cuda_graph = false,
bool grouped_all_reduce = false) {
auto j_optimizer = get_json(config, "optimizer");
auto embedding_name = get_value_from_json<std::string>(j_layers, "type");
auto bottom_name = get_value_from_json<std::string>(j_layers, "bottom");
auto top_name = get_value_from_json<std::string>(j_layers, "top");
auto& embed_wgrad_buff = (grouped_all_reduce) ?
std::dynamic_pointer_cast<GroupedExchangeWgrad<TypeFP>>(exchange_wgrad)->get_embed_wgrad_buffs() :
std::dynamic_pointer_cast<NetworkExchangeWgrad<TypeFP>>(exchange_wgrad)->get_embed_wgrad_buffs();
auto j_hparam = get_json(j_layers, "sparse_embedding_hparam");
size_t max_vocabulary_size_per_gpu = 0;
if (embedding_type == Embedding_t::DistributedSlotSparseEmbeddingHash) {
max_vocabulary_size_per_gpu =
get_value_from_json<size_t>(j_hparam, "max_vocabulary_size_per_gpu");
} else if (embedding_type == Embedding_t::LocalizedSlotSparseEmbeddingHash) {
if (has_key_(j_hparam, "max_vocabulary_size_per_gpu")) {
max_vocabulary_size_per_gpu =
get_value_from_json<size_t>(j_hparam, "max_vocabulary_size_per_gpu");
} else if (!has_key_(j_hparam, "slot_size_array")) {
CK_THROW_(Error_t::WrongInput,
"No max_vocabulary_size_per_gpu or slot_size_array in: " + embedding_name);
}
}
auto embedding_vec_size = get_value_from_json<size_t>(j_hparam, "embedding_vec_size");
auto combiner = get_value_from_json<int>(j_hparam, "combiner");
SparseInput<TypeKey> sparse_input;
if (!find_item_in_map(sparse_input, bottom_name, sparse_input_map)) {
CK_THROW_(Error_t::WrongInput, "Cannot find bottom");
}
OptParams<TypeFP> embedding_opt_params;
if (has_key_(j_layers, "optimizer")) {
embedding_opt_params = get_optimizer_param<TypeFP>(get_json(j_layers, "optimizer"));
} else {
embedding_opt_params = get_optimizer_param<TypeFP>(j_optimizer);
}
embedding_opt_params.scaler = scaler;
switch (embedding_type) {
case Embedding_t::DistributedSlotSparseEmbeddingHash: {
const SparseEmbeddingHashParams<TypeFP> embedding_params = {
batch_size,
batch_size_eval,
max_vocabulary_size_per_gpu,
{},
embedding_vec_size,
sparse_input.max_feature_num_per_sample,
sparse_input.slot_num,
combiner, // combiner: 0-sum, 1-mean
embedding_opt_params};
embeddings.emplace_back(new DistributedSlotSparseEmbeddingHash<TypeKey, TypeFP>(
sparse_input.train_row_offsets, sparse_input.train_values, sparse_input.train_nnz,
sparse_input.evaluate_row_offsets, sparse_input.evaluate_values,
sparse_input.evaluate_nnz, embedding_params, resource_manager));
break;
}
case Embedding_t::LocalizedSlotSparseEmbeddingHash: {
#ifndef NCCL_A2A
auto j_plan = get_json(j_layers, "plan_file");
std::string plan_file;
if (j_plan.is_array()) {
int num_nodes = j_plan.size();
plan_file = j_plan[resource_manager->get_process_id()].get<std::string>();
} else {
plan_file = get_value_from_json<std::string>(j_layers, "plan_file");
}
std::ifstream ifs(plan_file);
#else
std::string plan_file = "";
#endif
std::vector<size_t> slot_size_array;
if (has_key_(j_hparam, "slot_size_array")) {
auto slots = get_json(j_hparam, "slot_size_array");
for (auto slot : slots) {
slot_size_array.emplace_back(slot.get<size_t>());
}
}
const SparseEmbeddingHashParams<TypeFP> embedding_params = {
batch_size,
batch_size_eval,
max_vocabulary_size_per_gpu,
slot_size_array,
embedding_vec_size,
sparse_input.max_feature_num_per_sample,
sparse_input.slot_num,
combiner, // combiner: 0-sum, 1-mean
embedding_opt_params};
embeddings.emplace_back(new LocalizedSlotSparseEmbeddingHash<TypeKey, TypeFP>(
sparse_input.train_row_offsets, sparse_input.train_values, sparse_input.train_nnz,
sparse_input.evaluate_row_offsets, sparse_input.evaluate_values,
sparse_input.evaluate_nnz, embedding_params, plan_file, resource_manager));
break;
}
case Embedding_t::LocalizedSlotSparseEmbeddingOneHot: {
// 省略部分程式碼
break;
}
case Embedding_t::HybridSparseEmbedding: {
// 省略部分程式碼
break;
}
} // switch
// 這裡設定輸出
for (size_t i = 0; i < resource_manager->get_local_gpu_count(); i++) {
train_tensor_entries_list[i].push_back(
{top_name, (embeddings.back()->get_train_output_tensors())[i]});
evaluate_tensor_entries_list[i].push_back(
{top_name, (embeddings.back()->get_evaluate_output_tensors())[i]});
}
}
5.2.4 如何得到輸出
上面 create_embeddings 方法呼叫時候，train_tensor_entries_list 作為引數傳入。
在create_embeddings結尾處，會取出 embedding 的輸出，設定在 train_tensor_entries_list 之中。
// 如果是dense tensor，就會設定在這裡
for (size_t i = 0; i < resource_manager->get_local_gpu_count(); i++) {
train_tensor_entries_list[i].push_back(
{top_name, (embeddings.back()->get_train_output_tensors())[i]});
evaluate_tensor_entries_list[i].push_back(
{top_name, (embeddings.back()->get_evaluate_output_tensors())[i]});
}
輸出在哪裡？就是在 embedding_data.train_output_tensors_ 之中，後續我們會分析。
std::vector<TensorBag2> get_train_output_tensors() const override { \
std::vector<TensorBag2> bags; \
for (const auto& t : embedding_data.train_output_tensors_) { \
bags.push_back(t.shrink()); \
} \
return bags; \
} \
所以，對於 embedding，就是通過sparse_input_map 和 train_tensor_entries_list 構成了輸入，輸出資料流。
0xFF 參考
Using Neural Networks for Your Recommender System
Accelerating Embedding with the HugeCTR TensorFlow Embedding Plugin
NVIDIA Merlin HugeCTR 簡介：專用於推薦系統的培訓框架
https://web.eecs.umich.edu/~justincj/teaching/eecs442/notes/linear-backprop.html
稀疏矩陣儲存格式總結+儲存效率對比:COO,CSR,DIA,ELL,HYB
求通俗講解下tensorflow的embedding_lookup介面的意思？
【技術乾貨】聊聊在大廠推薦場景中embedding都是怎麼做的
ctr預估演算法對於序列特徵embedding可否做拼接，輸入MLP?與pooling
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/operations.html