为了提高大规模推荐系统的效率,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 性能的更多详细见解,您可以查看 研究论文。
Image source: Shutterstock