NVIDIA 的 EMBark 彻底改变大规模推荐系统训练

为了提高大规模推荐系统的效率，NVIDIA 推出了 EMBark，这是一种旨在优化深度学习推荐模型中的嵌入过程的新方法。根据 NVIDIA 的说法，推荐系统对于互联网行业至关重要，而有效地训练它们对许多公司来说是一个重大挑战。

推荐系统训练中的挑战

深度学习推荐模型（DLRMs）通常包含数十亿个 ID 特征，因此需要强大的训练解决方案。最近的 GPU 技术进步，例如 NVIDIA Merlin HugeCTR 和 TorchRec，通过利用 GPU 内存来处理大规模 ID 特征嵌入物改进了 DLRM 训练。然而，随着 GPU 数量的增加，嵌入期间的通信开销成为瓶颈，有时占据了总训练开销的一半以上。

EMBark 的创新方法

在 RecSys 2024 上发布的 EMBark 通过实施 3D 灵活分片策略和通信压缩技术来解决这些挑战，旨在在训练期间平衡负载并减少嵌入物的通信时间。EMBark 系统包括三个核心组件：嵌入集群、灵活的 3D 分片方案和分片规划器。

嵌入集群

这些集群将相似的特征分组，并应用定制的压缩策略，促进高效训练。EMBark 将集群分类为数据并行（DP）、基于减少（RB）和基于唯一（UB）类型，每种类型适用于不同的训练场景。

灵活的 3D 分片方案

这一创新方案允许对 GPU 之间的工作负载平衡进行精确控制，使用 3D 元组表示每个分片。这种灵活性解决了传统分片方法中存在的不平衡问题。

分片规划器

分片规划器使用贪婪搜索算法来确定最佳分片策略，基于硬件和嵌入配置来增强训练过程。

性能和评估

在 NVIDIA DGX H100 节点上测试的 EMBark 展示了显著的训练吞吐量改进。在各种 DLRM 模型中，EMBark 实现了平均 1.5 倍的训练速度提升，有些配置比传统方法快达 1.77 倍。

通过增强嵌入过程，EMBark 显著改进了大规模推荐系统模型的效率，为深度学习推荐系统设立了新的标准。有关 EMBark 性能的更多详细见解，您可以查看 研究论文。