如何使用NVIDIA Merlin推薦系統(tǒng)框架實現(xiàn)嵌入優(yōu)化

推薦系統(tǒng)是互聯(lián)網(wǎng)的經(jīng)濟引擎。很難想象任何其他類型的應用程序會對我們的日常數(shù)字生活產(chǎn)生更直接的影響：數(shù)以萬億計的項目被推薦給數(shù)十億人。推薦系統(tǒng)會在大量選項中過濾產(chǎn)品和服務，從而緩解大多數(shù)用戶面臨的選擇悖論。

隨著數(shù)據(jù)量的增加，深度學習（ DL ）推薦系統(tǒng)開始顯示與傳統(tǒng)的基于機器學習的方法相比的優(yōu)勢，例如梯度增強樹。為了給出一個具體的數(shù)據(jù)點， NVIDIA RAPIDS 。 AI 團隊與 DL 贏得了三場推薦比賽最近：

ACM RecSys2021 挑戰(zhàn)

SIGIR 和日期挑戰(zhàn)

ACM WSDM2021 Booking 。 com 挑戰(zhàn)賽

甚至在一年前 NVIDIA 數(shù)據(jù)科學家詢問為什么深度學習模型還沒有在推薦系統(tǒng)競賽中持續(xù)獲勝？時，這種情況也沒有持續(xù)發(fā)生。

嵌入在現(xiàn)代基于 DL 的推薦體系結(jié)構中起著關鍵作用，為數(shù)十億實體（用戶、產(chǎn)品及其特征）編碼單個信息。隨著數(shù)據(jù)量的增加，嵌入表的大小也隨之增加，現(xiàn)在這些表跨越多個 GB 到 TB 。在訓練這種類型的 DL 系統(tǒng)時存在著獨特的挑戰(zhàn)，其龐大的嵌入表和稀疏的訪問模式可能跨越多個 GPU 節(jié)點（如果不是節(jié)點的話）。

本文重點介紹 NVIDIA Merlin 推薦系統(tǒng)框架如何應對這些挑戰(zhàn)，并介紹了一種優(yōu)化的嵌入實現(xiàn)，其性能比其他框架的嵌入層高出 8 倍。這個優(yōu)化的實現(xiàn)也可以作為一個 TensorFlow 插件提供，它可以與 TensorFlow 無縫地工作，并作為 TensorFlow 本機嵌入層的方便替代品。

Embeddings

Embeddings 是一種機器學習技術，它將每個感興趣的對象（用戶、產(chǎn)品、類別等）表示為密集的數(shù)字向量。因此，嵌入表只是一種特定類型的鍵值存儲，鍵值是用于唯一標識對象的 ID ，值是實數(shù)向量。

Embeddings 是現(xiàn)代 DL 推薦系統(tǒng)中的一個關鍵構建塊，通常位于輸入層之后、“特征交互”和密集層之前。嵌入層是從數(shù)據(jù)和端到端訓練中學習的，就像深層神經(jīng)網(wǎng)絡的其他層一樣。正是嵌入層將 DL recommender 模型與其他類型的 DL 工作負載區(qū)分開來：它們?yōu)槟Ｐ吞峁┝舜罅康?a target="_blank">參數(shù)，但幾乎不需要計算，而計算密集型密集層的參數(shù)數(shù)量要少得多。

舉一個具體的例子：原始博大精深模型有幾個密集的層，大小為［1024 、 512 、 256］，因此只有幾百萬個參數(shù)，而其嵌入層可以有數(shù)十億個條目和數(shù)十億個參數(shù)。例如，這與 NLP 領域流行的 BERT 模型體系結(jié)構形成對比，其中嵌入層只有數(shù)萬個條目，總計數(shù)百萬個參數(shù)，但密集的前饋和注意層由數(shù)億個參數(shù)組成。這種差異還導致了另一個觀察：與其他類型的 DL 模型相比， DL recommender 網(wǎng)絡每字節(jié)輸入數(shù)據(jù)的計算量通常要小得多。

為什么優(yōu)化嵌入對推薦者工作流很重要

為了理解為什么優(yōu)化嵌入層和相關操作很重要，以下是培訓嵌入的挑戰(zhàn)：大小和訪問速度。

Size

隨著在線平臺和服務獲得數(shù)億甚至數(shù)十億用戶，以及提供的獨特產(chǎn)品數(shù)量達到數(shù)十億，嵌入表的規(guī)模不斷擴大也就不足為奇了。

據(jù)報道 Instagram 已被正在開發(fā)大小達到 10 TB 的推薦機型刪除。同樣，百度報告了廣告排名模型也達到了 10TB 的境界。在整個行業(yè)中，數(shù)百 GB 到 TB 的型號越來越流行，例如Pinterest 的 4-TB 模型和谷歌的 1 。 2 TB 模型。

很自然，在單個計算節(jié)點上擬合 TB 規(guī)模的模型是一個重大挑戰(zhàn)，更不用說在單個計算加速器上了。作為參考，最大的 NVIDIA A100 GPU 目前配備了 80 GB 的 HBM 。

訪問速度

訓練推薦系統(tǒng)本質(zhì)上是一項內(nèi)存帶寬密集型任務。這是因為每個訓練樣本或批通常在嵌入表中涉及少量實體。必須檢索這些條目以計算向前傳遞，然后在向后傳遞中更新。

CPU 主存儲器具有高容量，但帶寬有限，高端型號通常在幾十 GB / s 的范圍內(nèi)。另一方面， GPU 的內(nèi)存容量有限，但帶寬很高。 NVIDIA A100 80-GB GPU 提供 2 TB / s 的內(nèi)存帶寬。

解決方案

這些挑戰(zhàn)以不同的方式得到解決。例如，將整個嵌入表保留在主存上可以解決大小問題。然而，它通常會導致極低的訓練吞吐量，這往往與新數(shù)據(jù)的數(shù)量和速度相形見絀，從而使系統(tǒng)無法及時重新訓練。

或者，嵌入可以被仔細地分散在多個 GPU S 和多個節(jié)點上，僅被通信瓶頸所困擾，導致持續(xù)嚴重的 GPU – 在使用和訓練性能下的計算與純 CPU 訓練正好一致。

嵌入層是推薦系統(tǒng)的主要瓶頸之一。優(yōu)化嵌入層是解鎖 GPU 高計算吞吐量的關鍵。

在下一節(jié)中，我們將討論 NVIDIA Merlin HugeCTR 推薦框架如何通過使用 NVIDIA 技術，如 GPU 直接遠程直接內(nèi)存訪問（ RDMA ）、 NVIDIA 集體通信圖書館（ NCCL ）、 NVLink 和 NVSwitch ，解決大規(guī)模嵌入式的挑戰(zhàn)。它解鎖了 GPU 的高計算和高帶寬容量，同時解決了開箱即用、多 GPU 、多節(jié)點支持和模型并行性的內(nèi)存容量問題。

NVIDIA Merlin HugeCTR 嵌入概述

NVIDIA Merlin 解決了培訓大規(guī)模推薦系統(tǒng)的挑戰(zhàn)。它是一個端到端的推薦框架，可以加速推薦系統(tǒng)開發(fā)的所有階段，從數(shù)據(jù)預處理到培訓和推理。 NVIDIA Merlin HugeCTR 是一個開源的推薦系統(tǒng)，專用 DL 框架。在這篇文章中，我們關注 HugeCTR 的一個特定方面：嵌入優(yōu)化。

有兩種方法可以利用 HugeCTR 中的嵌入優(yōu)化工作：

將本機 NVIDIA Merlin HugeCTR 框架用于培訓和推理工作負載

使用 NVIDIA Merlin HugeCTR TensorFlow 插件，該插件旨在與 TensorFlow 無縫協(xié)作

原生 HugeCTR 嵌入優(yōu)化

為了克服嵌入挑戰(zhàn)并實現(xiàn)更快的訓練， HugeCTR 實現(xiàn)了自己的嵌入層，其中包括 GPU 加速哈希表、以節(jié)省內(nèi)存的方式實現(xiàn)的高效稀疏優(yōu)化器以及各種嵌入分布策略。它利用NCCL作為其內(nèi)部 GPU 通信原語。

哈希表實現(xiàn)基于 RAPIDS cuDF ，它是 GPU 數(shù)據(jù)幀庫，構成 NVIDIA 的 RAPIDS 數(shù)據(jù)科學平臺的一部分。 cuDF GPU 哈希表可以實現(xiàn)比 CPU 實現(xiàn)高達 35 倍的加速，例如來自線程構造塊（ TBB ）的并發(fā)_ hash _映射。有關更多信息，請參閱介紹 NVIDIA Merlin HugeCTR ：一個專門用于推薦系統(tǒng)的培訓框架。

考慮到可伸縮性， HugeCTR 默認支持嵌入層的模型并行性。嵌入表分布在可用的 GPU 和節(jié)點上。另一方面，密集層采用數(shù)據(jù)并行（圖 1 ）。

圖 1 。 HugeCTR 中的嵌入層并行性

騰訊推薦團隊是原生 HugeCTR 框架的首批采用者之一，大量使用了原生嵌入層。騰訊廣告與深度學習平臺負責人孔祥亭在最近的一次interview中表示，“ HugeCTR 作為一個推薦培訓框架，被整合到［騰訊］廣告推薦培訓系統(tǒng)中，以加快模型培訓的更新頻率，并可以培訓更多樣本以提高在線效果?！?/p>

HugeCTR TensorFlow 插件

NVIDIA Merlin 框架的所有組件都是開源的，旨在與更大的深度學習和數(shù)據(jù)科學生態(tài)系統(tǒng)進行互操作。我們的長期愿景是加速 GPU 上的推薦工作負載，而不管您喜歡哪種框架。 HugeCTR TensorFlow 嵌入插件是為了實現(xiàn)這一目標而創(chuàng)建的。

在高層次上， TensorFlow 嵌入插件是通過利用許多與本機 HugeCTR 嵌入層相同的嵌入優(yōu)化技術來設計的。特別是，這將是 GPU 哈希表和 NCCL ，位于引擎蓋下，用于 GPU 之間的通信。

HugeCTR 嵌入插件設計為方便無縫地與 TensorFlow 配合使用，作為 TensorFlow 本機嵌入層（如tf.nn.embedding_lookup和tf.nn.embedding_lookup_sparse）的替代品。它還提供了開箱即用的高級功能，例如在多個 GPU 上分布嵌入表的模型并行性。

NVIDIA Merlin HugeCTR TensorFlow 插件演練

下面是如何使用 TensorFlow 嵌入插件。完整示例可在 HugeCTR repository上找到，我們還提供了完整的基準測試筆記本用于再現(xiàn)性能數(shù)據(jù)。

訪問 HugeCTR 嵌入插件最方便的方式是使用 NGC NVIDIA Merlin TensorFlow 培訓碼頭工人形象，在該插件中預編譯和安裝 NVIDIA Merlin 框架的其他組件以及 TensorFlow 。最新的版本可以直接從 HugeCTR 存儲庫中提取，實時編譯和安裝。更新 TensorFlow 時，還必須為新安裝的 TensorFlow 版本重新編譯插件。

為了進行比較，這里介紹了如何使用本機 TensorFlow 嵌入層。首先，初始化 2D 數(shù)組變量以保存嵌入的值。然后，使用 tf 。 nn 。 embedding _ lookup 查找 ID 列表對應的嵌入值。

保存嵌入的值。然后，使用 tf . nn . embedding _ lookup 查找 ID 列表對應的嵌入值。

embedding_var = tf.Variable(initial_value=initial_value, dtype=tf.float32, name='embedding_variables')

@tf.function
def _train_step(inputs, labels):
    emb_vectors = tf.nn.embedding_lookup([self.embedding_var], inputs)
    ...

for i, (inputs, labels) in enumerate(dataset):
    _train_step(inputs)

同樣，也可以使用 HugeCTR 嵌入插件。首先，初始化嵌入層。接下來，該嵌入層用于查找 ID 列表的相應嵌入值。

import sparse_operation_kit as sok

emb_layer = sok.All2AllDenseEmbedding(max_vocabulary_size_per_gpu,
                                      embedding_vec_size,
                                      slot_num, nnz_per_slot)

@tf.function
def _train_step(inputs, labels):
    emb_vectors = emb_layer(inputs)
    ...

for i, (inputs, labels) in enumerate(dataset):
    _train_step(inputs)

HugeCTR 嵌入插件旨在與 TensorFlow 無縫協(xié)作，包括 Adam 和 sgd 等其他層和優(yōu)化器。在 TensorFlow v2 。 5 之前， Adam 優(yōu)化器是基于 CPU 的實現(xiàn)。

為了充分發(fā)揮 HugeCTR 嵌入插件的潛力，我們還在 sok 。 optimizers 。 adam 中提供了一個基于 GPU 的插件版本。從 TensorFlow v2 。 5 開始，標準 Adam 優(yōu)化器 tf 。 keras 。 optimizers 。 Adam （現(xiàn)在帶有 GPU 實現(xiàn)）可以以類似的精度和性能使用。

績效基準

在本節(jié)中，我們通過合成和實際用例展示 HugeCTR TensorFlow 嵌入插件的性能。

合成數(shù)據(jù)

在本例中，我們使用具有 100 個特征字段的合成數(shù)據(jù)集，每個都有 10 個查找，詞匯表大小為 8192 。推薦模型是一個具有六層的 MLP ，每層大小為 1024 。使用 TensorFlow 中的精確模型架構、優(yōu)化器和數(shù)據(jù)加載器，我們觀察到在 1x A100 GPU 上， HugeCTR 嵌入插件比本機 TensorFlow 嵌入層的平均迭代時間提高了 7 。 9 倍（圖 2 ）。

當從 1 到 4 個 A100 HugeCTR s 進行強縮放時，我們觀察到總加速比為 23 。 6 倍。默認情況下， GPU 嵌入插件提供了多重 GPU 縮放的好處。在引擎蓋下，嵌入插件以模型并行方式自動將與功能字段對應的表格分發(fā)到可用的 GPU 上。這與本機 TensorFlow 嵌入層形成對比，在本機嵌入層中，分布式模型并行多 GPU 訓練需要大量額外工作。 TensorFlow 分布策略 Mirrored 策略和 MultiWorkerMirrored 策略都是為進行數(shù)據(jù)并行同步訓練而設計的。

圖 2 。 HugeCTR TensorFlow 在合成數(shù)據(jù)集上 NVIDIA DGX A100 80GB 上嵌入插件的性能

真實用例：美團推薦系統(tǒng)

美團推薦系統(tǒng)團隊是首批成功采用 HugeCTR TensorFlow 插件的團隊之一。起初，團隊基于 CPU 優(yōu)化了他們的培訓框架，但隨著他們的模型變得越來越復雜，很難對培訓框架進行更深入的優(yōu)化?，F(xiàn)在，美團正致力于將 NVIDIA HugeCTR 集成到基于 A100 GPU 的培訓系統(tǒng)中。

美團高級技術專家黃軍（ Jun Huang ）表示：“一臺配備 8x A100 GPU s 的服務器可以替代基于 CPU 的培訓系統(tǒng)中的數(shù)百名員工。成本也大大降低。這是初步的優(yōu)化結(jié)果，未來仍有很大的優(yōu)化空間。”。

美團使用DIEN作為推薦模型。嵌入?yún)?shù)的總數(shù)是數(shù)百億，每個樣本中有數(shù)千個特征字段。由于輸入特征的范圍事先不是固定的和未知的，團隊使用哈希表在輸入嵌入層之前唯一地標識每個輸入特征。

使用 TensorFlow 中的精確模型架構、優(yōu)化器和數(shù)據(jù)加載器，我們觀察到在單個 A100 GPU 上， HugeCTR 嵌入插件實現(xiàn)了比原始 TensorFlow 嵌入 11 。 5 倍的加速。在弱縮放情況下， 8x A100 GPU s 上的迭代時間僅略微增加到 1x A100 GPU 的 1 。 17 倍（圖 3 ）。

圖 3 。美團數(shù)據(jù)上 NVIDIA DGX A100 80GB 上的 HugeCTR TensorFlow 嵌入插件性能

結(jié)論

HugeCTR TensorFlow 嵌入插件今天可以從 HugeCTR GitHub 存儲庫以及 NGC NVIDIA Merlin TensorFlow 容器獲得。如果您是 TensorFlow 用戶，希望構建和部署具有大型嵌入表的大規(guī)模推薦系統(tǒng)， HugeCTR TensorFlow 插件可以輕松替換 TensorFlow 嵌入查找層。

關于作者

Vinh Nguyen 是一位深度學習的工程師和數(shù)據(jù)科學家，發(fā)表了 50 多篇科學文章，引文超過 2500 篇。在 NVIDIA ，他的工作涉及廣泛的深度學習和人工智能應用，包括語音、語言和視覺處理以及推薦系統(tǒng)。

Ann Spencer 是 Merlin 的高級產(chǎn)品營銷經(jīng)理。在加入 NVIDIA 之前，她曾在多家數(shù)據(jù)公司擔任產(chǎn)品和研究職務。 2012-2014 年，她還是 O ‘ Reilly Media 的數(shù)據(jù)編輯，專注于數(shù)據(jù)工程和數(shù)據(jù)科學。

Joey Wang 是 NVIDIA Merlin HugeCTR 和 RnnEngine 的作者和開發(fā)經(jīng)理，也是 NVIDIA 亞太地區(qū)發(fā)展技術解決方案團隊的經(jīng)理。他在 NVIDIA 工作了七年。自 2014 年以來， Joey 主要從事深度學習 GPU 加速工作。在此之前，喬伊在 HPC 從事智能視頻分析和粒子模擬。

Jianbing Dong 是 DevTech 工程師， HugeCTR 嵌入插件的作者。在 NVIDIA ，他的工作重點是 TensorFlow 相關的優(yōu)化。

審核編輯：郭婷