如何在TensorFlow2中高效培訓(xùn)1130億參數(shù)推薦系統(tǒng)

深度學(xué)習(xí)推薦系統(tǒng)通常使用大型嵌入表。很難將它們放入 GPU 內(nèi)存中。

這篇文章向你展示了如何結(jié)合使用模型并行和數(shù)據(jù)并行訓(xùn)練范例來解決這個記憶問題，從而更快地訓(xùn)練大型深度學(xué)習(xí)推薦系統(tǒng)。我分享了我的團(tuán)隊在 TensorFlow 2 中高效培訓(xùn) 1130 億參數(shù)推薦系統(tǒng)所采取的步驟，該模型的所有嵌入的總大小為 421 GiB 。

通過在 GPU 和 CPU 之間拆分模型和嵌入，我的團(tuán)隊實現(xiàn)了 43 倍的加速。然而，將嵌入分布到多個 GPU 上，帶來了令人難以置信的 672 倍的加速。這種多 GPU 方法實現(xiàn)了顯著的加速，使您能夠在幾分鐘內(nèi)而不是幾天內(nèi)訓(xùn)練大型推薦系統(tǒng)。

您可以使用 NVIDIA 深度學(xué)習(xí)示例 GitHub 存儲庫 中提供的代碼自己復(fù)制這些結(jié)果。

嵌入層的模型并行訓(xùn)練

在數(shù)據(jù)并行訓(xùn)練中，每個 GPU 存儲模型的相同副本，但在不同的數(shù)據(jù)上訓(xùn)練。這對于許多深度學(xué)習(xí)應(yīng)用程序來說都很方便，因為它易于實現(xiàn)，并且通信開銷相對較低。然而，這種模式要求神經(jīng)網(wǎng)絡(luò)的權(quán)重適合單個設(shè)備。

如果模型大小大于設(shè)備內(nèi)存，一種方法是將模型分成子部分，并在不同的 GPU 上訓(xùn)練每個子部分。這被稱為模型并行訓(xùn)練。

表的每一行對應(yīng)于要映射到密集表示的輸入變量的值。表中的每一列表示輸出空間的不同維度，表示所有向量中一個值的切片。因為一個典型的深度學(xué)習(xí)推薦程序會吸收多個分類特征，所以它需要多個嵌入表。

對于具有多個大型嵌入的推薦程序，有三種實現(xiàn)模型并行性的方法：

Table-wise split——每個嵌入表完全放在一個設(shè)備上；每個設(shè)備只包含所有嵌入的一個子集。（圖 1 ）

Column-wise split–每個 GPU 包含每個嵌入表中的一個子集列。（圖 2 ）

Row-wise split–每個 GPU 保存每個嵌入表中的行子集。

由于負(fù)載平衡問題，行分割比其他兩個選項更難實現(xiàn)。在本文中，我將重點介紹表拆分和列拆分?；旌虾推ヅ涠喾N方法是一個可行的選擇，但為了簡單起見，我不會在本文中集中討論這一點。

圖 1 。按表拆分模式是指每個 GPU 存儲所有嵌入表的子集

圖 2 。按列拆分模式是指每個設(shè)備存儲來自每個嵌入表的列的子集

這些方法之間有一些關(guān)鍵區(qū)別（表 1 ）。簡言之，按表拆分模式更易于使用，而且可能更快，具體取決于具體的工作負(fù)載。

一個缺點是它不支持嵌入跨越多個 GPU 的表。相比之下，按列拆分模式支持嵌入跨多個 GPU 的表，但速度可能會稍慢，尤其是對于窄表。

表 1 。表拆分和列拆分模式之間的比較。

高效訓(xùn)練推薦系統(tǒng)的混合并行方法

典型的推薦程序在嵌入后運行算術(shù)密集型層，如線性或點積。處理模型這一部分的一種幼稚方法是將嵌入查找的結(jié)果收集到單個 GPU 上，并在此 GPU 上運行這些密集層。然而，這是非常低效的，因為在這段時間內(nèi)沒有使用用于保存嵌入的另一個 GPU 。

更好的方法是使用所有 GPU 通過數(shù)據(jù)并行運行密集層。這可以通過按批量大小拆分嵌入查找的結(jié)果來實現(xiàn)。也就是說，對于 N 和八 GPU 的全局批量，每個 GPU 只處理 N / 8 個訓(xùn)練樣本。實際上，這意味著密集層以數(shù)據(jù)并行模式運行。

由于這種方法結(jié)合了嵌入的模型并行性和多層感知器（ MLP ）的數(shù)據(jù)并行性，因此被稱為混合并行訓(xùn)練（圖 3 ）。

圖 3 。訓(xùn)練大型推薦系統(tǒng)的通用混合并行方法

Horovod all-to-all

從模型并行到數(shù)據(jù)并行范式需要一個多 GPU 集體通信操作：全部對全部。

All to All 是一種靈活的集體通信原語，可在每對 GPU 之間交換數(shù)據(jù)。這是必需的，因為在嵌入查找階段結(jié)束時，每個 GPU 都保存所有樣本的查找結(jié)果。但是，僅適用于表的子集（用于按表拆分）或列的子集（用于按列拆分）。

由于 all-to-all 操作會在 GPU 之間洗牌數(shù)據(jù)，因此需要注意的是，每個 GPU 都會保存所有表的所有列的嵌入查找結(jié)果，但只保存樣本子集的嵌入查找結(jié)果。例如，對于一個 8 GPU 場景，本地批量大小畢竟是之前的 8 倍。

通信由 Horovod 庫的 hvd.alltoall 函數(shù)處理。在引擎蓋下，霍洛伍德稱 NCCL 實施 為了獲得最佳性能。如果你的系統(tǒng)上有 NVLink ，它也會利用它。

TensorFlow 2 中的混合并行訓(xùn)練示例

在本節(jié)中，我將描述一種用于 TensorFlow 2 中訓(xùn)練的 1130 億參數(shù)推薦系統(tǒng)的混合并行訓(xùn)練方法。完整的源代碼可以在 NVIDIA 深度學(xué)習(xí)示例庫 中找到。

深度學(xué)習(xí)推薦模型的體系結(jié)構(gòu)

對于這個例子，我使用 DLRM 體系結(jié)構(gòu)（圖 4 ）。 DLRM 是研究論文 面向個性化和推薦系統(tǒng)的深度學(xué)習(xí)推薦模型 中首次介紹的一類推薦模型。我之所以選擇它，是因為 MLPerf 基準(zhǔn)測試使用了更小版本的 DLRM ，因此，它是演示推薦系統(tǒng)性能的當(dāng)前行業(yè)標(biāo)準(zhǔn)。

DLRM 同時使用分類和數(shù)字功能。分類特征被輸入到嵌入層中，而數(shù)字特征則由一個小的 MLP 子網(wǎng)絡(luò)處理。

然后將這些層的結(jié)果輸入點交互層和另一個 MLP 。然后使用二元交叉熵?fù)p失函數(shù)通過反向傳播對模型進(jìn)行訓(xùn)練，并根據(jù)隨機梯度下降（ SGD ）方法更新權(quán)重。

圖 4 。 DLRM 體系結(jié)構(gòu)圖。

修改以支持寬深度模型

雖然我選擇在本例中使用 DLRM 體系結(jié)構(gòu)，但也可以支持相關(guān)模型（如 Wide & Deep ）。這需要進(jìn)行以下修改：

添加 wide & Deep 的“ wide ”部分，并在純數(shù)據(jù)并行模式下運行它，完全繞過 all to all 。

為寬部分添加第二個優(yōu)化器。

在深部，移除底部 MLP ，并將數(shù)字特征直接傳遞到頂部 MLP 。

移除點交互層。

同步文件夾

DLRM 可以在由數(shù)字和分類特征組成的任何表格數(shù)據(jù)集上進(jìn)行訓(xùn)練。在本例中，我使用 Criteo 的 TB 點擊日志數(shù)據(jù)集 ，因為它是最大的公開點擊率數(shù)據(jù)集。

該數(shù)據(jù)集由 26 個分類變量和 13 個數(shù)值變量組成。在未經(jīng)處理的數(shù)據(jù)中，獨特類別的總數(shù)為 8.82 億，其中 2.92 億是在最大的特征中發(fā)現(xiàn)的。

遵循 MLPerf 推薦基準(zhǔn)，對嵌入使用單精度，每個特征的嵌入維度為 128 。這意味著參數(shù)總數(shù)為 882M × 128 = 1130 億。所有 26 個表的總大小為 1130 億× 4 字節(jié)/ 230= 421 GiB ，最大表為 139.6 GiB 。因為最大的表不適合單個 GPU ，所以必須使用按列拆分模式將表分片，并將每個表分布到多個 GPU 中。

從理論上講，您可以只對超過單個 GPU 內(nèi)存的少數(shù)表執(zhí)行此操作，并對其余的表使用按表拆分。然而，這將不必要地使代碼復(fù)雜化，而沒有任何明顯的好處。因此，對所有表使用按列拆分模式。

性能優(yōu)化

為了提高訓(xùn)練速度，我的團(tuán)隊實施了以下性能優(yōu)化，如代碼所示。這些是可以應(yīng)用于其他深度學(xué)習(xí)推薦系統(tǒng)以及其他深度學(xué)習(xí)框架的通用策略。

自動混合精度

混合精度是計算方法中不同數(shù)值精度的組合使用。有關(guān)如何啟用它的更多信息，請參閱 TensorFlow 核心文檔中的 Mixed precision 。與 A100 上默認(rèn)的 TF32 精度相比，該模型使用混合精度使其速度提高了 23% 。

相同寬度的融合嵌入表

當(dāng)多個嵌入表具有相同的向量大小時——這是 DLRM 中使用embedding_dim=128的情況——它們可以沿零軸連接。這允許對一個大表執(zhí)行單個查找，而不是對許多較小的表執(zhí)行多個查找。

啟動一個大內(nèi)核而不是多個小內(nèi)核要高效得多。在本例中，將表連接起來可使訓(xùn)練速度提高 39% 。

XLA

我的團(tuán)隊使用 TensorFlow 加速線性代數(shù)（ XLA ）編譯器來提高性能。對于這個特定的用例，應(yīng)用 XLA 比不使用它產(chǎn)生 3.36X 的加速。這個值是在打開所有其他優(yōu)化的情況下實現(xiàn)的： AMP 、串聯(lián)嵌入等等。

廣播數(shù)據(jù)加載器

在每個 GPU 上運行每個嵌入表的一部分意味著每個 GPU 必須訪問每個訓(xùn)練樣本的每個特性。在每個過程中分別加載和解析所有這些輸入數(shù)據(jù)效率低下，可能會導(dǎo)致嚴(yán)重的瓶頸。我通過只在第一個 worker 上加載輸入數(shù)據(jù)并通過 NVLink 將其廣播給其他 worker 來解決這個問題。這提供了 32% 的加速。

把這一切放在一起

圖 5 顯示了具有八個 GPU 的混合并行 DLRM 的設(shè)備放置示例。該圖顯示 GPU 0 和 7 。為了簡單起見，它只顯示分類功能 0 和 25 。

圖 5 。具有 1130 億個參數(shù)的混合并行 DLRM 的簡化圖。

替代方法：將大型嵌入存儲在 CPU 上

存儲大型嵌入矩陣的一個簡單替代方法是將它們放入主機內(nèi)存中。小型嵌入表和計算密集型層仍然可以放置在 GPU 上，以獲得最佳性能。雖然簡單得多，但與將所有變量保留在 GPU 上相比，這種方法也較慢。

這有兩個根本原因：

嵌入查找是一種內(nèi)存受限的操作。 CPU 內(nèi)存比 GPU 內(nèi)存慢得多。雙插槽 AMD Epyc 7742 的總內(nèi)存帶寬為 409.6 GB / s ，而單插槽 A100-80GB GPU 的總內(nèi)存帶寬為 2 TB / s ，而 8 個 A100-80GB GPU 的總內(nèi)存帶寬為 16 TB / s 。

GPU 之間的數(shù)據(jù)交換速度明顯快于 CPU 和 GPU 之間的數(shù)據(jù)交換速度。這是因為將 CPU 連接到 GPU 之間的 PCIe 鏈路可能會成為瓶頸。

當(dāng)使用 CPU 存儲嵌入時， CPU 和 GPU 之間的傳輸必須首先通過提供 31.5 GB / s 帶寬的 PCIe 接口。相反，在混合并行范例中，嵌入查找的結(jié)果通過 GPU 之間的 NVSwitch 結(jié)構(gòu)進(jìn)行傳輸。 DGX A100 采用第二代 NVSwitch 技術(shù)，支持每秒 600 GB 的峰值 GPU 到 – GPU 通信。

盡管速度有所放緩，但這種替代方法仍然比僅在 CPU 上運行整個網(wǎng)絡(luò)快得多。

基準(zhǔn)結(jié)果

下表顯示了訓(xùn)練 113B 參數(shù) DLRM 模型的基準(zhǔn)測試結(jié)果。它只比較了三種硬件設(shè)置： CPU ，一種使用 CPU 內(nèi)存的單一 GPU 用于最大的嵌入表，以及一種使用完整 DGX A100-80GB 的混合并行方法。

表 2 。比較 1130 億參數(shù)深度學(xué)習(xí)推薦模型（ DLRM ）的 CPU 和 GPU 訓(xùn)練吞吐量。

比較前兩行，你可以看到用一個 A100 GPU 來補充兩個 CPU 可以使吞吐量增加 43 倍。之所以會出現(xiàn)這種情況，是因為 GPU 非常適合運行計算密集型線性層和適合其 80-GB 內(nèi)存的較小嵌入層。

此外，使用八個 GPU 的完整 DGX A100 比在單個 A100 GPU 上訓(xùn)練快 15.5 倍。 DGX A100 使您能夠?qū)⒄麄€型號安裝到 GPU 內(nèi)存中，并消除了昂貴的設(shè)備到主機和主機到設(shè)備傳輸?shù)男枰?/p>

總的來說， DGX A100 解決這項任務(wù)的速度是雙插座 CPU 系統(tǒng)的 672 倍。

結(jié)論

在這篇文章中，我介紹了使用混合并行來訓(xùn)練大型推薦系統(tǒng)的想法。測試結(jié)果表明， DGX A100 是在 TensorFlow 2 中訓(xùn)練參數(shù)超過 1000 億的推薦系統(tǒng)的極好工具。它在雙插槽 CPU 上實現(xiàn)了 672 倍的加速。

高內(nèi)存帶寬和快速的 GPU 到 – GPU 通信使快速培訓(xùn)推薦人成為可能。因此，與僅使用 CPU 服務(wù)器相比，您的培訓(xùn)時間更短。這降低了培訓(xùn)成本，同時為從業(yè)者提供了更快的實驗。

關(guān)于作者

Tomasz Grel 是一名深度學(xué)習(xí)工程師。在NVIDIA ，他專注于確保眾多推薦系統(tǒng)的質(zhì)量和執(zhí)行速度，包括 NCF 、 VAE-CF 和 DLRM 。

審核編輯：郭婷