推荐系统是互联网行业的核心，而高效地训练这些系统对于各大公司来说是一个关键问题。大多数推荐系统是深度学习推荐模型（DLRMs），包含数十亿甚至数百亿个 ID 特征。图 1 示出了一个典型的结构。
近年来， NVIDIA Merlin HugeCTR 和 TorchRec 等 GPU 解决方案通过在 GPU 上存储大规模 ID 特征嵌入并对其进行并行处理，显著加速了 DLRM 的训练。与 CPU 解决方案相比，使用 GPU 内存带宽可实现显著改进。
与此同时，随着训练集群中使用的 GPU 数量增加（从 8 个 GPU 增加到 128 个 GPU），我们发现嵌入的通信开销在总训练开销中占较大比例。在一些大规模训练场景中（例如在 16 个节点上），它甚至超过了一半（51%）。
这主要有两个原因：
1、随着集群中 GPU 数量的增加，每个节点上的嵌入表数量逐渐减少，导致不同节点之间的负载不平衡，并降低训练效率。
2. 与节点内带宽相比，节点间带宽要小得多，因此嵌入模型并行的通信时间更长。
为了帮助行业用户更好地理解和解决这些问题，NVIDIA HugeCTR 团队在 RecSys 2024 上展示了 EMBark。EMBark 支持 3D 灵活的分片策略，并结合通信压缩策略来微调大规模集群深度推荐模型训练中的负载不平衡问题，并减少嵌入所需的通信时间。 相关代码 和 论文 均为开源。
EMBark 简介
EMBark 可提高不同集群配置下 DLRM 训练中嵌入的性能，并加速训练吞吐量。它使用 NVIDIA Merlin HugeCTR 开源推荐系统框架实施，但这些技术也可以应用于其他机器学习框架。
EMBark 包含三个关键组件：嵌入集群、灵活的 3D 分片方案和分片规划器。图 3 显示了 EMBark 的整体架构。
嵌入集群
嵌入集群可高效训练嵌入，方法是将嵌入按类似功能分组，并为每个集群应用自定义压缩策略。这些集群包括数据分布器、嵌入存储和嵌入运算符，它们共同工作，将特征 ID 转换为嵌入向量。
嵌入集群有三种类型：数据并行（DP）、基于归约（RB）和基于唯一（UB）。每种方法在训练期间都使用不同的通信方法，并且适用于不同的嵌入。
- DP 集群： DP 集群不会压缩通信，使通信变得简单高效，但它们会在每个 GPU 上复制嵌入表，因此仅适用于小型表。
- RB 集群 ：使用归约运算，对于使用池化运算压缩多热输入表格具有显著效果。
- UB 集群： 仅发送唯一向量，对于处理具有明显访问热点的嵌入表是有益的。
灵活的 3D 分片方案
灵活的 3D 分片方案可以解决 RB 集群中的工作负载不平衡问题。与固定分片策略（如逐行、逐表或逐列）不同，EMBark 使用 3D 元组 (i, j, k) 表示每个分片，其中 i 表示表格索引，j 表示行分片索引，k 表示列分片索引。这种方法允许每个嵌入跨任意数量的 GPU 进行分片，从而提供对工作负载平衡的灵活性和精确控制。
分片规划器
为了找到最佳的分片策略，EMBark 提供了一个分片规划器，这是一种基于成本的贪婪搜索算法，可以根据硬件规格和嵌入配置确定最佳的分片策略。
评估
利用嵌入集群、灵活的 3D 分片方案和高级分片规划器，EMBark 可显著提升训练性能。为了评估其有效性，我们在一组 NVIDIA DGX H100 节点上进行了实验，每个节点配备 8 个 NVIDIA H100 GPU (总计 640 GB HBM，每个节点的带宽为 24TB/s)。
在每个节点内，所有 GPU 都通过 NVLink（900 GB/s 双向）互联，节点间通信使用 InfiniBand（8x400Gbps）。
| DLRM-DCNv2 | T180 | T200 | T510 | T7 | |
| # parameters | 48B | 70B | 100B | 110B | 470B |
| FLOPS per sample | 96M | 308M | 387M | 1,015M | 25M |
| # embedding tables | 26 | 180 | 200 | 510 | 7 |
| Table dimensions (d) | 128 | {32, 128} | 128 | 128 | {32, 128} |
| Average hotness (h) | ~10 | ~80 | ~20 | ~5 | ~20 |
| Achieved QPS | ~75M | ~11M | ~15M | ~8M | ~74M |
为了证明 EMBark 可以高效地训练任何规模的 DLRM 模型，我们使用 MLPerf DLRM-DCNv2 模型进行了测试，并生成了多个具有更大嵌入表和不同属性的合成模型（请参阅表 1）。我们的训练数据集显示幂律偏差为 a = 1.2。
基准使用串行内核执行顺序、固定的表格行分片策略以及所有 RB 集群。实验先后使用了三种优化：重叠、更灵活的分片策略和更好的集群配置。在四个典型的 DLRM 变体（DLRM-DCNv2、T180、T200 和 T510）中，EMBark 平均实现了 1.5 倍的端到端训练吞吐量加速，比基准快达 1.77 倍。
您可以参阅 论文 ，了解更多详细的实验结果和相关分析。
使用 EMBark 加快 DLRM 训练速度
EMBark 解决了大规模推荐系统模型训练中漫长嵌入过程的挑战。通过支持 3D 灵活的分片策略，结合不同的通信压缩策略，可以在大规模集群上微调深度推荐模型训练中的负载不平衡问题，并减少嵌入所需的通信时间。
这提高了大规模推荐系统模型的训练效率。在四个代表性的 DLRM 变体（DLRM-DCNv2、T180、T200 和 T510）中，EMBark 平均实现了 1.5 倍的端到端训练吞吐量加速，比基准快 1.77 倍。
我们还在积极探索嵌入卸载相关技术和优化 torch 相关工作，并期待在未来分享我们的进展。如果您对此领域感兴趣， 请联系我们 获取更多信息。
