NVIDIA Merlin 助力陌陌推荐系统训练吞吐提升最高 12 倍

一、方案背景与核心挑战

挚文集团旗下陌陌是基于地理位置的移动视频社交应用，是中国领先的开放式社交平台之一。随着用户规模增加和业务发展，推荐算法准确度要求提高，导致模型复杂性和训练样本量显著增加。核心挑战：陌陌原有方案本质是基于 PS-Worker 的 CPU + GPU 混合训练方案，可支持大规模稀疏参数训练，但在性能上难以满足业务日益增长的需求，对单次模型训练速度和新模型算法探索效率都构成更大挑战，亟需优化训练速度以加快算法迭代。

二、解决方案名称与架构

解决方案名称：基于 NVIDIA GPU 与 Merlin 的陌陌推荐系统训练加速方案

核心技术：NVIDIA 团队与陌陌推荐系统团队深度合作，使用 NVIDIA GPU 和 Merlin 软件解决方案替代原有方案，核心组件为 Merlin TensorFlow Plugin（即 Sparse Operation Kit，SOK）和 HierarchicalKV（HKV）。在不影响模型效果的前提下，模型整体吞吐提升 5 倍以上，再结合通信和 IO 进一步优化后，极限情况下可提升 12 倍吞吐。

三、核心技术组件

SOK 模型并行：NVIDIA Merlin HugeCTR 可高效利用 GPU 进行推荐系统训练，Merlin TensorFlow Plugin 将 HugeCTR 的高级特性封装为 TensorFlow 可直接调用的形式。SOK 以数据并行方式接收输入数据，将稀疏参数以模型并行方式分布在多个 GPU 上，稠密参数以数据并行方式分布，内部实现"数据并行—模型并行—数据并行"的转换流程。其使用方式尽可能与原有 TensorFlow 算子对齐，减少代码修改，并针对 embedding vector 的拷贝与 combiner 做了高度优化与内核融合。

HKV 层次化动态嵌入存储：Merlin HierarchicalKV 是面向推荐系统训练设计的 KV 加速库，为兼容大模型训练支持层次化动态 Embedding 存储（CPU+GPU）、灵活的 eviction 机制和丰富 API，目前已集成入 SOK 协同加速 Embedding 计算。在陌陌实际应用中，动态 Embedding 特性大大简化了连续训练时需人工控制显存中 embedding tab 大小的问题。

端到端性能优化：陌陌推荐团队在 SOK + HKV 架构基础上结合业务特点进一步优化，包括梯度合并减少梯度计算通信开销、并行特征数据读取与转换以及特征数据预取到 GPU 提速特征 IO、使用 XLA 编译优化融合 kernel 以减少 kernel launch 时间、设置 GPU 亲和等，使整体性能提升达到 12 倍。

四、实施成果

双方团队通过 SOK 和 HKV 对原方案进行深度优化后，成功帮助陌陌提升 12 倍训练效率，极大降低了模型训练成本和新模型算法尝试成本，整体方案已上线并全面支持陌陌推荐系统模型训练。近期双方还针对基于 Transformer 的推荐模型开展性能优化，将 XLA+AMP+半精度 allreduce 应用到该模型后端到端性能实现 50% 加速；进一步优化 multi-head-attention 热点并集成 Flash-Attention 后，预计整体加速比可达 3 倍，同时显存使用量下降约 70%。

五、未来展望

陌陌与 NVIDIA 将继续在推荐系统训练和推理等方面持续合作，推进 GPU 和 AI 软件加速计算在陌陌的全面落地。