NVIDIA Merlin分布式嵌入使快速、TB級(jí)推薦培訓(xùn)變得簡(jiǎn)單

嵌入在深度學(xué)習(xí)推薦模型中起著關(guān)鍵作用。它們用于將數(shù)據(jù)中的編碼分類輸入映射到可由數(shù)學(xué)層或多層感知器（ MLP ）處理的數(shù)值。

嵌入通常構(gòu)成深度學(xué)習(xí)推薦模型中的大部分參數(shù)，并且可以相當(dāng)大，甚至達(dá)到 TB 級(jí)。在訓(xùn)練期間，很難將它們放入單個(gè) GPU 的內(nèi)存中。

因此，現(xiàn)代推薦者可能需要模型并行和數(shù)據(jù)并行分布式訓(xùn)練方法的組合，以實(shí)現(xiàn)合理的訓(xùn)練時(shí)間和可用 GPU 計(jì)算的最佳利用。

NVIDIA Merlin 分布式嵌入 ，在 TensorFlow 2 中，一個(gè)用于訓(xùn)練大型基于嵌入的（例如，推薦者）模型的庫使您只需幾行代碼即可輕松完成。

背景

通過 GPU 上的數(shù)據(jù)并行分布式訓(xùn)練，在每個(gè) GPU 工作人員上復(fù)制整個(gè)模型。在訓(xùn)練過程中，一批數(shù)據(jù)在多個(gè) GPU 中分割，每個(gè)設(shè)備獨(dú)立操作其自己的數(shù)據(jù)碎片。

這允許將計(jì)算擴(kuò)展到更大批量的更高數(shù)據(jù)量。在反向傳播期間計(jì)算的梯度使用減少操作（例如， horovod.tensorflow.allreduce ） 用于同步參數(shù)更新。

通過模型并行分布式訓(xùn)練，模型參數(shù)在不同工作人員之間進(jìn)行分割。這是一種更適合分發(fā)大型嵌入表的方法。訓(xùn)練需要使用全對(duì)全通信原語（例如， horovod.tensorflow.alltoall ） 使得工人可以訪問不在其分區(qū)中的參數(shù)。

在之前的相關(guān)文章中， 在 TensorFlow 2 中使用 100B +參數(shù)在 DGX A100 上訓(xùn)練推薦系統(tǒng) ， Tomasz 討論了如何為 1130 億參數(shù)分配嵌入 DLRM 跨多個(gè) NVIDIA GPU 的模型有助于在僅 CPU 的解決方案上實(shí)現(xiàn) 672 倍的加速。這一重大改進(jìn)可能會(huì)將訓(xùn)練時(shí)間從幾天縮短到幾分鐘！這是通過模型并行分布嵌入表和通過數(shù)據(jù)并行執(zhí)行小得多的數(shù)學(xué)密集型 MLP 層計(jì)算來實(shí)現(xiàn)的。

與將嵌入存儲(chǔ)在 CPU 內(nèi)存中相比，這種混合方法使您能夠使用 GPU 內(nèi)存的高內(nèi)存帶寬進(jìn)行內(nèi)存綁定嵌入查找。它還使用幾個(gè) GPU 設(shè)備中的計(jì)算能力加速 MLP 層。作為參考 NVIDIA A100-80GB GPU 具有帶寬超過 2 TB / s 的 80 GB HBM2 存儲(chǔ)器）。

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

嵌入表可以分為“表方式”（例如，嵌入表 0 和 N ）、“列方式”（例如嵌入表 2 ）或“行方式”?？缢?GPU 復(fù)制 MLP 層。數(shù)字特征可以直接輸入 MLP 層，并且在圖中未示出。

然而，實(shí)現(xiàn)這種復(fù)雜的混合并行訓(xùn)練方法并不簡(jiǎn)單，需要領(lǐng)域?qū)＜以O(shè)計(jì)幾百行低級(jí)代碼來開發(fā)和優(yōu)化訓(xùn)練。

為了使其更廣泛地使用 NVIDIA Merlin 分布式嵌入 該庫提供了一個(gè)易于使用的包裝器，只需三行 Python 代碼即可在 TensorFlow 2 中民主化模型并行性。它提供了一個(gè)可伸縮的模型并行包裝器 分發(fā)嵌入表 除了一些 高效嵌入操作 這涵蓋并擴(kuò)展了 TensorFlow 的嵌入功能。下面是它如何實(shí)現(xiàn)混合并行。

分布式并行模型

NVIDIA Merlin 分布式嵌入提供了 distributed_embeddings.dist_model_parallel 單元。它有助于在多個(gè) GPU 工作者之間分發(fā)嵌入，而無需任何復(fù)雜的代碼來處理與原語的跨工作者通信，如 all2all 下面的代碼示例顯示了此 API 的用法：

import dist_model_parallel as dmp

class MyEmbeddingModel(tf.keras.Model):
  def  __init__(self, table_sizes):
    ...
    self.embedding_layers = [tf.keras.layers.Embedding(input_dim, output_dim) for input_dim, output_dim in table_sizes]
    # 1. Add this line to wrap list of embedding layers used in the model
    self.embedding_layers = dmp.DistributedEmbedding(self.embedding_layers)
  def call(self, inputs):
    # embedding_outputs = [e(i) for e, i in zip(self.embedding_layers, inputs)]
    embedding_outputs = self.embedding_layers(inputs)
    ...

要使用 Horovod 以數(shù)據(jù)并行方式運(yùn)行密集層，請(qǐng)?zhí)鎿QHorovod’sDistributed GradientTape和broadcast方法及其在分布式嵌入中的等效。以下示例直接取自 Horovod 文檔，并進(jìn)行了相應(yīng)修改。

@tf.function
def training_step(inputs, labels, first_batch):
  with tf.GradientTape() as tape:
    probs = model(inputs)
    loss_value = loss(labels, probs)

  # 2. Change Horovod Gradient Tape to dmp tape
  # tape = hvd.DistributedGradientTape(tape)
  tape = dmp.DistributedGradientTape(tape)
  grads = tape.gradient(loss_value, model.trainable_variables)
  opt.apply_gradients(zip(grads, model.trainable_variables))

  if first_batch:
    # 3. Change Horovod broadcast_variables to dmp's
    # hvd.broadcast_variables(model.variables, root_rank=0)
    dmp.broadcast_variables(model.variables, root_rank=0)
  return loss_value

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

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

性能

為了證明使用 NVIDIA Merlin 分布式嵌入的好處，我們展示了在 Criteo 1TB 數(shù)據(jù)集上訓(xùn)練的 DLRM 模型的基準(zhǔn)測(cè)試，以及各種具有多達(dá) 3 個(gè) TiB 嵌入表大小的合成模型。

Criteo 數(shù)據(jù)集上的 DLRM 基準(zhǔn)

基準(zhǔn)測(cè)試表明，我們使用更簡(jiǎn)單的 API 保持了類似于專家工程代碼的性能。這個(gè) NVIDIA 深度學(xué)習(xí)示例 DLRM 使用 TensorFlow 2 的代碼現(xiàn)在也已更新，以利用 NVIDIA Merlin 分布式嵌入的混合并行訓(xùn)練。更多信息，請(qǐng)參閱我們之前的文章， 在 TensorFlow 2 中使用 100B +參數(shù)在 DGX A100 上訓(xùn)練推薦系統(tǒng) 。

這個(gè) 基準(zhǔn) 自述部分提供了對(duì)性能數(shù)字的更多了解。

具有 1130 億個(gè)參數(shù)（ 421 個(gè) GiB 模型大?。┑?DLRM 模型在 Criteo TB 點(diǎn)擊日志 數(shù)據(jù)集，三種不同的硬件設(shè)置：

僅 CPU 的解決方案。

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

使用 NVIDIA DGX A100-80GB 和 8 GPU 的混合并行解決方案。這利用了 NVIDIA Merlin 分布式嵌入提供的模型并行包裝器和嵌入 API 。

我們觀察到， DGX-A100 上的分布式嵌入解決方案比僅使用 CPU 的解決方案提供了驚人的 683 倍的加速！我們還注意到，與單一 GPU 解決方案相比，性能有了顯著改善。這是因?yàn)樵?GPU 內(nèi)存中保留所有嵌入消除了通過 CPU-GPU 接口嵌入查找的開銷。

綜合模型基準(zhǔn)

為了進(jìn)一步演示解決方案的可伸縮性，我們創(chuàng)建了不同大小的合成 DLRM 模型（表 2 ）。

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

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

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

結(jié)論

在這篇文章中，我們介紹了 NVIDIA Merlin 分布式嵌入庫，僅需幾行代碼即可在 NVIDIA GPU 上實(shí)現(xiàn)基于嵌入的深度學(xué)習(xí)模型的可擴(kuò)展和高效模型并行訓(xùn)練。

關(guān)于作者

Shashank Verma 是 NVIDIA 的一名深入學(xué)習(xí)的技術(shù)營(yíng)銷工程師。他負(fù)責(zé)開發(fā)和展示各種深度學(xué)習(xí)框架中以開發(fā)人員為中心的內(nèi)容。他從威斯康星大學(xué)麥迪遜分校獲得電氣工程碩士學(xué)位，在那里他專注于計(jì)算機(jī)視覺、數(shù)據(jù)科學(xué)的安全方面和 HPC 。

Wenwen Gao 是 NVIDIA Merlin 的高級(jí)產(chǎn)品經(jīng)理，擁有 Amazon 和其他技術(shù)公司的產(chǎn)品管理經(jīng)驗(yàn)，專注于個(gè)性化和推薦。她擁有多倫多大學(xué)計(jì)算機(jī)科學(xué)學(xué)士學(xué)位和麻省理工學(xué)院斯隆管理學(xué)院工商管理碩士學(xué)位。

Hao Wu 是 NVIDIA 的高級(jí) GPU 計(jì)算架構(gòu)師。他在完成博士學(xué)位后于 2011 年加入 NVIDIA 計(jì)算架構(gòu)組。在中國(guó)科學(xué)院。近年來， Hao 的技術(shù)重點(diǎn)是將低精度應(yīng)用于深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練和推理。

Deyu Fu 是 NVIDIA 深度學(xué)習(xí)框架團(tuán)隊(duì)的高級(jí)開發(fā)技術(shù)工程師，負(fù)責(zé)加速軟件堆棧 CUDA 內(nèi)核、數(shù)學(xué)、通信、框架和模型的 DL 培訓(xùn)工作。他最近專注于 NVIDIA Merlin 分布式嵌入和推薦系統(tǒng)。

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

審核編輯：郭婷