Embedding在深度學習推薦模型中起著關鍵作用。它們被用于將輸入數據中的離散特征映射到向量，以便下游的神經網絡進行處理。Embedding 通常構成深度學習推薦模型中的大部分參數，大小可以達到 TB 級。在訓練期間，很難將它們放入單個 GPU 的內存中。因此，現代推薦系統可能需要模型并行和數據并行的分布式訓練方法組合，以最佳利用GPU計算資源來實現最好的訓練性能。
NVIDIA Merlin Distributed-Embeddings ，可以方便TensorFlow 2 用戶用短短幾行代碼輕松完成大規模的推薦模型訓練。
背景
在數據并行分布式訓練中，整個模型被復制到每個 GPU 上。在訓練過程中，一批輸入數據在多個 GPU 中分割，每張卡獨立處理其自己的數據分片，從而允許計算擴展到更大批量的數據。在反向傳播期間，計算的梯度通過reduction算子（例如， horovod.tensorflow.allreduce ) 來同步更新多個GPU間的參數。
另一方面，模型并行分布式訓練中，模型參數被分割到多個GPU上。這種方法更適合分布存儲大型embedding。訓練中，每個GPU 通過alltoall通信算子（例如，?? ) 訪問不在本機中的參數。
在之前的相關文章中，?用TensorFlow 2 在DGX A100上訓練100B +參數的推薦系統?， Tomasz 討論了如何將1130 億參數的DLRM?模型中的embedding分布到多個 NVIDIA GPU 進行訓練，并相比純CPU 的方案實現 672 倍的性能提升。這一重大突破可以將訓練時間從幾天縮短到幾分鐘！這是通過模型并行 embedding層和數據并行MLP層來實現的。和CPU方案相比，這種混合并行的方法能夠有效利用GPU 的高內存帶寬加速內存受限的embedding查找，并同時利用多個 GPU 的算力加速 MLP 層。作為參考，?NVIDIA A100-80GB GPU?具有超過 2 TB / s 的帶寬和 80 GB HBM2 存儲）。
圖 1.用于訓練大型推薦系統的通用“混合并行”方法
embedding表可以按表為分割單位（圖中表 0 和 N ），按”列“分割（圖中表 2），或者按”行”分割。MLP 層跨所有 GPU 復制，而數字特征則可以直接輸入 MLP 層。
然而，實現這種復雜的混合并行訓練方法并不簡單，需要領域內專家設計數百行底層代碼來開發和優化。為了使其更普適，??NVIDIA Merlin Distributed-Embeddings 提供了一些易于使用的TensorFlow 2的封裝，讓所有人都只需三行 Python 代碼即可輕松實現模型并行。它提供了一些涵蓋并拓展原生TensorFlow功能的高性能embedding查找算子。在此基礎上，它提供了一個可規模化的模型并行封裝函數，幫助用戶自動將embedding分布于多個GPU上。下面將展示它如何實現混合并行。
分布式模型并行
NVIDIA Merlin Distributed-Embeddings 提供了?distributed_embeddings.dist_model_parallel?模塊。它有助于在多個 GPU 之間分布embedding而無需任何復雜的代碼來處理跨GPU間的通信（如 ?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 以數據并行方式運行MLP層，請將 Horovod的?Distributed GradientTape? 和?broadcast?方法替換成 NVIDIA Merlin Distributed-Embeddings 里同等的API。以下示例直接取自 Horovod 文檔，并進行了相對應修改。
@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
通過這些微小的改變，您就可以使用混合并行訓練了！
我們還提供了以下完整示例：?使用 Criteo 1TB 點擊日志數據訓練 DLRM 模型?以及?擴展到 22.8 TiB的合成數據?模型。
性能
為了展示 NVIDIA Merlin Distributed-Embeddings 的性能，我們在 Criteo 1TB 數據集DLRM 模型和最高達到 3 TiB embedding的合成模型上進行了模型訓練的基準測試。
Criteo 數據集上的 DLRM 模型基準測試
測試表明，我們使用更簡單的 API 取得了近似于專家代碼的性能。NVIDIA深度學習DLRM TensorFlow 2示例代碼現已更新為使用 NVIDIA Merlin Distributed-Embeddings 進行分布式混合并行訓練，更多信息請參閱我們之前的文章，?用TensorFlow 2 在DGX A100上訓練100B +參數的推薦系統. README中的?基準測試?部分提供了對性能結果的更多詳述。
?
我們對1130 億個參數（ 421 個 GiB 大小）的 DLRM 模型在?Criteo TB 點擊日志?數據集上用三種不同的硬件設置進行了訓練：
- 僅 CPU 的解決方案。
- 單 GPU 解決方案，其中 CPU 內存用于存儲最大的embedding表。
- 使用 NVIDIA DGX A100-80GB 的 8 GPU 的混合并行解決方案。此方案利用了 NVIDIA Merlin Distributed-Embeddings 里提供的模型并行api和embedding API 。
|
Hardware |
Description |
Training Throughput (samples/second) |
Speedup over CPU |
|
2 x AMD EPYC 7742 |
Both MLP layers and embeddings on CPU |
17.7k |
1x |
|
1 x A100-80GB; 2 x AMD EPYC 7742 |
Large embeddings on CPU, everything else on GPU |
768k |
43x |
|
DGX A100 (8xA100-80GB) |
Hybrid parallel with NVIDIA Merlin?Distributed-Embeddings,?whole model on GPU |
12.1M |
683x |
表 1.各種設置的訓練吞吐量和加速
我們觀察到， DGX-A100 上的 NVIDIA Merlin Distributed-Embeddings 方案比僅使用 CPU 的解決方案提供了驚人的 683 倍的加速！我們還注意到與單GPU 方案相比，混合并行的性能也有顯著提升。這是因為在 GPU 顯存中存儲所有embedding避免了通過 CPU-GPU 接口查找embedding的開銷。
合成模型基準測試
為了進一步演示方案的可規模化，我們創建了不同大小的合成數據以及對應的 DLRM 模型（表 2 ）。有關模型生成方法和訓練腳本的更多信息，請參見 GitHub NVIDIA-Merlin/distributed-embeddings?代碼庫。
|
Model |
Total number of embedding tables |
Total embedding size (GiB) |
|
Tiny |
55 |
4.2 |
|
Small |
107 |
26.3 |
|
Medium |
311 |
206.2 |
|
Large |
612 |
773.8 |
|
Jumbo |
1,022 |
3,109.5 |
表 2.合成模型尺寸
每個合成模型使用一個或多個 DGX-A100-80GB 節點進行訓練，全局數據batch大小為 65536 ，并使用 Adagrad 優化器。從表 3 中可以看出，? NVIDIA Merlin Distributed-Embeddings 可以在數百個 GPU 上輕松訓練 TB 級模型。
|
Model |
Training step time (ms) |
||||
|
1 GPU |
8 GPU |
16 GPU |
32 GPU |
128 GPU |
|
|
Tiny |
17.6 |
3.6 |
3.2 |
? |
? |
|
Small |
57.8 |
14.0 |
11.6 |
7.4 |
? |
|
Medium |
? |
64.4 |
44.9 |
31.1 |
17.2 |
|
Large |
? |
? |
? |
65.0 |
33.4 |
|
Jumbo |
? |
? |
? |
? |
102.3 |
表 3.各種硬件配置下合成模型的訓練步長時間（ ms ）
另一方面，與傳統的數據并行相比，即使對于可以容納在單個 GPU 中的模型，多GPU分布式模型并行仍然提供了顯著加速。表 4 顯示了上述Tiny模型在DGX A100-80GB 上的性能對比。
|
Solution |
Training step time (ms) |
|||
|
1 GPU |
2 GPU |
4 GPU |
8 GPU |
|
|
NVIDIA Merlin Distributed Embeddings?Model Parallel |
17.7 |
11.6 |
6.4 |
4.2 |
|
Native TensorFlow?Data Parallel |
19.9 |
20.2 |
21.2 |
22.3 |
表 4. Tiny模型（ 4.2GiB ）的訓練步長時間（ ms ）比較 NVIDIA Merlin Distributed-Embeddings 模型并行和原生 TensorFlow 數據并行
本實驗使用了 65536 的全局批量和 Adagrad 優化器。
結論
在這篇文章中，我們介紹了? NVIDIA Merlin Distributed-Embeddings，僅需幾行代碼即可在 NVIDIA GPU 上實現基于 embedding的深度學習模型，并進行可規模化，高效率地模型并行訓練。歡迎嘗試以下?使用合成數據的可擴展訓練示例?和?基于 Criteo 數據訓練 DLRM 模型?示例.
?
