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

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

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

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

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

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

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

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

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

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

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

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

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

由于負(fù)載平衡問(wèn)題，行分割比其他兩個(gè)選項(xiàng)更難實(shí)現(xiàn)。在本文中，我將重點(diǎn)介紹表拆分和列拆分。混合和匹配多種方法是一個(gè)可行的選擇，但為了簡(jiǎn)單起見(jiàn)，我不會(huì)在本文中集中討論這一點(diǎn)。

圖 1 。按表拆分模式是指每個(gè) GPU 存儲(chǔ)所有嵌入表的子集

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

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

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

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

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

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

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

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

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

Horovod all-to-all

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

All to All 是一種靈活的集體通信原語(yǔ)，可在每對(duì) GPU 之間交換數(shù)據(jù)。這是必需的，因?yàn)樵谇度氩檎译A段結(jié)束時(shí)，每個(gè) GPU 都保存所有樣本的查找結(jié)果。但是，僅適用于表的子集（用于按表拆分）或列的子集（用于按列拆分）。

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

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

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

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

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

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

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

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

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

修改以支持寬深度模型

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

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

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

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

移除點(diǎn)交互層。

同步文件夾

DLRM 可以在由數(shù)字和分類(lèi)特征組成的任何表格數(shù)據(jù)集上進(jìn)行訓(xùn)練。在本例中，我使用 Criteo 的 TB 點(diǎn)擊日志數(shù)據(jù)集 ，因?yàn)樗亲畲蟮墓_(kāi)點(diǎn)擊率數(shù)據(jù)集。

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

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

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

性能優(yōu)化

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

自動(dòng)混合精度

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

相同寬度的融合嵌入表

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

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

XLA

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

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

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

把這一切放在一起

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

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

替代方法：將大型嵌入存儲(chǔ)在 CPU 上

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

這有兩個(gè)根本原因：

嵌入查找是一種內(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 個(gè) A100-80GB GPU 的總內(nèi)存帶寬為 16 TB / s 。

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

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

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

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

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

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

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

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

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

結(jié)論

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

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

關(guān)于作者

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

審核編輯：郭婷