[源碼解析] 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