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

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

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

ACM RecSys2021 挑戰(zhàn)

SIGIR 和日期挑戰(zhàn)

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

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

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

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

Embeddings

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

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

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

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

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

Size

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

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

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

訪問(wèn)速度

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

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

解決方案

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

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

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

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

NVIDIA Merlin HugeCTR 嵌入概述

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

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

將本機(jī) NVIDIA Merlin HugeCTR 框架用于培訓(xùn)和推理工作負(fù)載

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

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

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

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

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

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

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

HugeCTR TensorFlow 插件

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

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

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

NVIDIA Merlin HugeCTR TensorFlow 插件演練

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

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

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

保存嵌入的值。然后，使用 tf . nn . embedding _ lookup 查找 ID 列表對(duì)應(yīng)的嵌入值。

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 嵌入插件。首先，初始化嵌入層。接下來(lái)，該嵌入層用于查找 ID 列表的相應(yīng)嵌入值。

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 無(wú)縫協(xié)作，包括 Adam 和 sgd 等其他層和優(yōu)化器。在 TensorFlow v2 。 5 之前， Adam 優(yōu)化器是基于 CPU 的實(shí)現(xiàn)。

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

績(jī)效基準(zhǔn)

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

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

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

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

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

真實(shí)用例：美團(tuán)推薦系統(tǒng)

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

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

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

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

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

結(jié)論

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

關(guān)于作者

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

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

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

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

審核編輯：郭婷