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

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

背景

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

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

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

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

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

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

分布式模型并行

NVIDIA Merlin Distributed-Embeddings 提供了

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

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

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

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

性能

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

Criteo 數(shù)據(jù)集上的 DLRM 模型基準測試

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

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

僅 CPU 的解決方案。

單 GPU 解決方案，其中 CPU 內(nèi)存用于存儲最大的 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 方案相比，混合并行的性能也有顯著提升。這是因為在 GPU 顯存中存儲所有 embedding 避免了通過 CPU-GPU 接口查找 embedding 的開銷。

合成模型基準測試

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

表 2. 合成模型尺寸

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

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

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

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

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

結(jié)論

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