EMBark：面向大规模推荐系统嵌入训练的优化方法 NEWS DETAIL

当深度学习推荐模型（DLRM）的 ID 特征达到数十亿甚至数百亿规模时，嵌入表的存储、访问和跨 GPU 通信会成为训练扩展的关键限制。EMBark 是 NVIDIA HugeCTR 团队在 RecSys 2024 展示的一种嵌入训练优化方法，针对大规模 GPU 集群中的嵌入通信开销和负载不平衡，提供嵌入集群、灵活 3D 分片和分片规划器三项能力。它以 NVIDIA Merlin HugeCTR 开源推荐系统框架实现，但其方法可用于评估其他机器学习框架中的类似问题。

大规模 DLRM 训练面临什么问题

典型 DLRM 同时包含稠密特征与大量 ID 特征，后者通常需要通过嵌入表转换为向量。NVIDIA Merlin HugeCTR 和 TorchRec 等 GPU 方案可在 GPU 上存储和并行处理大规模 ID 特征嵌入，利用 GPU 内存带宽改善训练过程。

不过，集群从较少 GPU 扩展到更多 GPU 后，嵌入模型并行带来的通信占比可能上升。源资料指出，在一个 16 节点的大规模训练场景中，嵌入通信开销曾超过总训练开销的一半，达到 51%。其原因包括：每节点可承载的嵌入表分布变化可能导致节点间负载不均；同时，节点间带宽通常低于节点内带宽，使跨节点嵌入通信耗时更长。

EMBark 的三项核心能力

嵌入集群将功能相近的嵌入分组，并按组应用相应的通信或压缩策略。集群由数据分布器、嵌入存储和嵌入运算符组成，用于将特征 ID 转换为嵌入向量。

• 数据并行（DP）集群：不压缩通信，并在每张 GPU 上复制嵌入表，通信路径相对直接，适合小型表。
• 基于归约（RB）集群：采用归约运算，面向可通过池化操作压缩多热输入的表。
• 基于唯一（UB）集群：仅发送唯一向量，适合访问热点较明显的嵌入表。

灵活 3D 分片方案主要用于处理 RB 集群的工作负载不平衡。EMBark 以 (i, j, k) 描述分片：i 为表索引，j 为行分片索引，k 为列分片索引。相较于固定的逐行、逐表或逐列策略，这种表示可使一个嵌入跨任意数量 GPU 分片，为负载分配提供更细的控制粒度。

分片规划器采用基于成本的贪婪搜索算法，根据硬件规格和嵌入配置确定分片策略。它的作用不是替代部署验证，而是在表规模、访问模式、GPU 数量和互连条件已知时，为分片决策提供系统化起点。

适合评估的训练场景

EMBark 更适合评估具有以下特征的推荐训练任务：嵌入表规模大、训练需要跨多个节点扩展、嵌入通信已成为明显瓶颈，或现有固定分片方式造成 GPU 或节点负载不均。对于小型嵌入表，DP 集群的复制方式可能更易实现；对于多热输入且可池化压缩的表，可重点评估 RB 集群；对于访问分布存在明显热点的表，则可分析 UB 集群是否与实际访问模式匹配。

源资料中的实验环境为 NVIDIA DGX H100 节点：每节点配置 8 个 NVIDIA H100 GPU、总计 640 GB HBM，节点内通过 NVLink 互联，双向带宽为 900 GB/s；节点间使用 InfiniBand，配置为 8x400Gbps。该环境可说明评估所依赖的硬件条件，但不应直接外推为其他 GPU、网络拓扑、嵌入规模或数据分布下的性能结论。

建议的评估路径与边界

1. 梳理嵌入表数量、表容量、多热特征比例、访问热点和池化操作，确认通信与负载问题来自哪些表。
2. 记录节点内与节点间互连配置、GPU 数量、可用显存及框架版本，将这些条件作为分片规划输入。
3. 分别比较 DP、RB 与 UB 集群在目标数据集上的通信量、负载分布和端到端训练吞吐量，而非只观察单一算子时间。
4. 针对 RB 集群测试 3D 分片与既有固定分片策略，检查长尾 GPU、节点间流量及训练稳定性。
5. 在完整训练配置下进行项目测试，并以实际模型收敛、资源占用和端到端吞吐量决定是否采用。

源资料说明相关代码和论文为开源，但未提供具体版本、适配范围或可复现配置。实施前应查阅带日期的官方 NVIDIA Merlin HugeCTR 文档、对应代码仓库说明及完整软件物料清单，核实框架兼容性、依赖项和适用限制。

常见问题

EMBark 是否只适用于 NVIDIA Merlin HugeCTR？

源资料明确说明 EMBark 使用 NVIDIA Merlin HugeCTR 实现，同时指出其技术可应用于其他机器学习框架。是否能够直接迁移到某一具体框架，仍需根据该框架的嵌入算子、通信机制、分片接口和版本文档进行验证。

3D 分片一定会提升所有推荐模型的训练性能吗？

不能据此作出保证。3D 分片用于提高 RB 集群中的负载平衡灵活性，其收益取决于嵌入表结构、输入分布、GPU 数量、节点内外带宽和现有分片策略。应在目标模型和目标集群上进行端到端测试。

结论

EMBark 为大规模 DLRM 嵌入训练提供了按嵌入特征选择通信方式、以 3D 方式细化分片、并由成本模型辅助规划的路径。对于通信开销高或跨节点负载不均的推荐训练集群，它值得纳入技术评估；最终部署决策应以官方文档、完整配置和项目实测结果为依据。

围绕“EMBark：面向大规模推荐系统嵌入训练的优化方法”继续了解 采购与选型问答。