Embedding 在深度學(xué)習(xí)推薦模型中起著關(guān)鍵作用。它們被用于將輸入數(shù)據(jù)中的離散特征映射到向量，以便下游的神經(jīng)網(wǎng)絡(luò)進(jìn)行處理。Embedding 通常構(gòu)成深度學(xué)習(xí)推薦模型中的大部分參數(shù)，大小可以達(dá)到 TB 級(jí)。在訓(xùn)練期間，很難將它們放入單個(gè) GPU 的內(nèi)存中。因此，現(xiàn)代推薦系統(tǒng)可能需要模型并行和數(shù)據(jù)并行的分布式訓(xùn)練方法組合，以最佳利用 GPU 計(jì)算資源來(lái)實(shí)現(xiàn)最好的訓(xùn)練性能。

NVIDIA Merlin Distributed-Embeddings，可以方便TensorFlow 2 用戶用短短幾行代碼輕松完成大規(guī)模的推薦模型訓(xùn)練。

背景

在數(shù)據(jù)并行分布式訓(xùn)練中，整個(gè)模型被復(fù)制到每個(gè) GPU 上。在訓(xùn)練過(guò)程中，一批輸入數(shù)據(jù)在多個(gè) GPU 中分割，每張卡獨(dú)立處理其自己的數(shù)據(jù)分片，從而允許計(jì)算擴(kuò)展到更大批量的數(shù)據(jù)。在反向傳播期間，計(jì)算的梯度通過(guò) reduction 算子（例如， horovod.tensorflow.allreduce ） 來(lái)同步更新多個(gè) GPU 間的參數(shù)。

另一方面，模型并行分布式訓(xùn)練中，模型參數(shù)被分割到多個(gè) GPU 上。這種方法更適合分布存儲(chǔ)大型 embedding。訓(xùn)練中，每個(gè) GPU 通過(guò) alltoall 通信算子（例如， horovod.tensorflow.alltoall） 訪問(wèn)不在本機(jī)中的參數(shù)。

在之前的相關(guān)文章中， 用 TensorFlow 2 在 DGX A100 上訓(xùn)練 100B + 參數(shù)的推薦系統(tǒng) ， Tomasz 討論了如何將 1130 億參數(shù)的 DLRM 模型中的 embedding 分布到多個(gè) NVIDIA GPU 進(jìn)行訓(xùn)練，并相比純 CPU 的方案實(shí)現(xiàn) 672 倍的性能提升。這一重大突破可以將訓(xùn)練時(shí)間從幾天縮短到幾分鐘！這是通過(guò)模型并行 embedding 層和數(shù)據(jù)并行 MLP 層來(lái)實(shí)現(xiàn)的。和 CPU 方案相比，這種混合并行的方法能夠有效利用 GPU 的高內(nèi)存帶寬加速內(nèi)存受限的 embedding 查找，并同時(shí)利用多個(gè) GPU 的算力加速 MLP 層。作為參考， NVIDIA A100-80GB GPU 具有超過(guò) 2 TB / s 的帶寬和 80 GB HBM2 存儲(chǔ)）。

圖 1. 用于訓(xùn)練大型推薦系統(tǒng)的通用“混合并行”方法

embedding 表可以按表為分割單位（圖中表 0 和 N ），按“列”分割（圖中表 2），或者按”行”分割。MLP 層跨所有 GPU 復(fù)制，而數(shù)字特征則可以直接輸入 MLP 層。

然而，實(shí)現(xiàn)這種復(fù)雜的混合并行訓(xùn)練方法并不簡(jiǎn)單，需要領(lǐng)域內(nèi)專家設(shè)計(jì)數(shù)百行底層代碼來(lái)開發(fā)和優(yōu)化。為了使其更普適，NVIDIA Merlin Distributed-Embeddings 提供了一些易于使用的 TensorFlow 2 的封裝，讓所有人都只需三行 Python 代碼即可輕松實(shí)現(xiàn)模型并行。它提供了一些涵蓋并拓展原生 TensorFlow 功能的高性能 embedding 查找算子。在此基礎(chǔ)上，它提供了一個(gè)可規(guī)模化的模型并行封裝函數(shù)，幫助用戶自動(dòng)將 embedding 分布于多個(gè) GPU 上。下面將展示它如何實(shí)現(xiàn)混合并行。

分布式模型并行

NVIDIA Merlin Distributed-Embeddings 提供了

distributed_embeddings.dist_model_parallel 模塊。它有助于在多個(gè) GPU 之間分布embedding而無(wú)需任何復(fù)雜的代碼來(lái)處理跨GPU間的通信（如 all2all ）。下面的代碼示例顯示了此 API 的用法：

要使用 Horovod 以數(shù)據(jù)并行方式運(yùn)行 MLP 層，請(qǐng)將 Horovod的 Distributed GradientTape 和 broadcast 方法替換成 NVIDIA Merlin Distributed-Embeddings 里同等的 API。以下示例直接取自 Horovod 文檔，并進(jìn)行了相對(duì)應(yīng)修改。

通過(guò)這些微小的改變，您就可以使用混合并行訓(xùn)練了！

我們還提供了以下完整示例： 使用 Criteo 1TB 點(diǎn)擊日志數(shù)據(jù)訓(xùn)練 DLRM 模型以及擴(kuò)展到 22.8 TiB 的合成數(shù)據(jù)模型。

性能

為了展示 NVIDIA Merlin Distributed-Embeddings 的性能，我們?cè)?Criteo 1TB 數(shù)據(jù)集 DLRM 模型和最高達(dá)到 3 TiB embedding 的合成模型上進(jìn)行了模型訓(xùn)練的基準(zhǔn)測(cè)試。

Criteo 數(shù)據(jù)集上的 DLRM 模型基準(zhǔn)測(cè)試

測(cè)試表明，我們使用更簡(jiǎn)單的 API 取得了近似于專家代碼的性能。NVIDIA 深度學(xué)習(xí) DLRM TensorFlow 2 示例代碼現(xiàn)已更新為使用 NVIDIA Merlin Distributed-Embeddings 進(jìn)行分布式混合并行訓(xùn)練，更多信息請(qǐng)參閱我們之前的文章， 用 TensorFlow 2 在 DGX A100 上訓(xùn)練 100B + 參數(shù)的推薦系統(tǒng)。README 中的基準(zhǔn)測(cè)試部分提供了對(duì)性能結(jié)果的更多詳述。

我們對(duì) 1130 億個(gè)參數(shù)（ 421 個(gè) GiB 大?。┑?DLRM 模型在 Criteo TB 點(diǎn)擊日志數(shù)據(jù)集上用三種不同的硬件設(shè)置進(jìn)行了訓(xùn)練：

僅 CPU 的解決方案。

單 GPU 解決方案，其中 CPU 內(nèi)存用于存儲(chǔ)最大的 embedding 表。

使用 NVIDIA DGX A100-80GB 的 8 GPU 的混合并行解決方案。此方案利用了 NVIDIA Merlin Distributed-Embeddings 里提供的模型并行 api 和 embedding API 。

表 1. 各種設(shè)置的培訓(xùn)吞吐量和加速

我們觀察到， DGX-A100 上的 NVIDIA Merlin Distributed-Embeddings 方案比僅使用 CPU 的解決方案提供了驚人的 683 倍的加速！我們還注意到與單 GPU 方案相比，混合并行的性能也有顯著提升。這是因?yàn)樵?GPU 顯存中存儲(chǔ)所有 embedding 避免了通過(guò) CPU-GPU 接口查找 embedding 的開銷。

合成模型基準(zhǔn)測(cè)試

為了進(jìn)一步演示方案的可規(guī)?；覀儎?chuàng)建了不同大小的合成數(shù)據(jù)以及對(duì)應(yīng)的 DLRM 模型（表 2 ）。有關(guān)模型生成方法和訓(xùn)練腳本的更多信息，請(qǐng)參見(jiàn) GitHub NVIDIA-Merlin/distributed-embeddings 代碼庫(kù)。

表 2. 合成模型尺寸

每個(gè)合成模型使用一個(gè)或多個(gè) DGX-A100-80GB 節(jié)點(diǎn)進(jìn)行訓(xùn)練，全局?jǐn)?shù)據(jù) batch 大小為 65536 ，并使用 Adagrad 優(yōu)化器。從表 3 中可以看出， NVIDIA Merlin Distributed-Embeddings 可以在數(shù)百個(gè) GPU 上輕松訓(xùn)練 TB 級(jí)模型。

表 3. 各種硬件配置下合成模型的訓(xùn)練步長(zhǎng)時(shí)間（ ms ）

另一方面，與傳統(tǒng)的數(shù)據(jù)并行相比，即使對(duì)于可以容納在單個(gè) GPU 中的模型，多 GPU 分布式模型并行仍然提供了顯著加速。表 4 顯示了上述 Tiny 模型在 DGX A100-80GB 上的性能對(duì)比。

表 4. Tiny模型（ 4.2GiB ）的訓(xùn)練步長(zhǎng)時(shí)間（ ms ）比較 NVIDIA Merlin Distributed-Embeddings 模型并行和原生 TensorFlow 數(shù)據(jù)并行

本實(shí)驗(yàn)使用了 65536 的全局批量和 Adagrad 優(yōu)化器。

結(jié)論

在這篇文章中，我們介紹了 NVIDIA Merlin Distributed-Embeddings，僅需幾行代碼即可在 NVIDIA GPU 上實(shí)現(xiàn)基于 embedding 的深度學(xué)習(xí)模型，并進(jìn)行可規(guī)模化，高效率地模型并行訓(xùn)練。歡迎嘗試以下使用合成數(shù)據(jù)的可擴(kuò)展訓(xùn)練示例和基于 Criteo 數(shù)據(jù)訓(xùn)練 DLRM 模型示例。