通过 Merlin 大幅提升大规模深度多目标精排模型训练性能
本案例中,NVIDIA 团队与陌陌推荐系统团队深度合作,共同使用 NVIDIA GPU 和 Merlin 软件解决方案替代其原有解决方案。
通过使用 Merlin TensorFlow Plugin (即 Sparse Operation Kit,SOK) 和 HierarchicalKV(HKV),相较于原方案在相同规模模型和 GPU 下,显著提高了陌陌大规模深度多目标精排模型的训练性能。在不影响模型效果的前提下,模型整体吞吐提升了 5 倍以上,再结合通信和 IO 等进一步优化后,极限情况下可以提升 12 倍吞吐。
客户简介
挚文集团于 2011 年成立,2014 年 12 月 11 日在美国纳斯达克交易所挂牌上市(NASDAQ: MOMO),拥有陌陌、探探等多款手机应用,以及电影制作发行、节目制作等多元业务。陌陌是挚文集团于 2011 年 8 月推出的一款基于地理位置的移动视频社交应用,是中国领先的开放式社交平台之一。
训练速度面临挑战,需有效提升算法迭代
陌陌的原始解决方案本质是基于 PS-Worker 的 CPU + GPU 混合训练方案,可支持大规模稀疏参数的训练。然而,随着用户规模的增加和业务的发展,对于推荐算法的准确度也有了更高的要求。这导致模型的复杂性和训练样本量显著增加,对单次模型训练速度和新模型算法探索效率都有更大的挑战。尽管原方案在功能上支持了大规模稀疏参数的训练,但在性能上难以满足业务日益增长的需求。因此,陌陌亟需对训练速度进行优化,加快算法迭代,以提高业务效果。
SOK 和 HKV 为推荐系统提升性能与灵活性
NVIDIA Merlin HugeCTR 是 NVIDIA 推出的可以高效利用 GPU 来进行推荐系统训练的解决方案,为了使它能直接被其他 DL 用户,比如 TensorFlow 所直接使用,NVIDIA 开发了 Merlin TensorFlow Plugin (以下简称 SOK),将 HugeCTR 中的高级特性封装为 TensorFlow 可直接调用的形式,从而帮助用户在 TensorFlow 中直接使用 HugeCTR 中的高级特性来加速他们的推荐系统。
Merlin TensorFlow Plugin 以数据并行的方式接收输入数据,将稀疏参数以模型并行的方式分布在多个 GPU 上,将稠密参数以数据并行的方式分布在多个 GPU 上,内部实现“数据并行-模型并行-数据并行”的转换流程。整个使用方式上尽可能的与原有 TensorFlow 算子对齐,减少对用户已有的代码的修改,以更方便、快捷地在多个 GPU 上进行扩展。此外,SOK 针对 embedding vector 的拷贝和 combiner 进行了高度优化和内核融合,使整个 lookup 的前后向过程拥有更好的性能。
Merlin HierarchicalKV (以下简称HKV)是 Merlin 下的针对于推荐系统训练设计的 KV 加速库。为兼容大模型训练支持了层次化动态 Embedding 存储(CPU+GPU),灵活的 eviction(淘汰) 机制和丰富的 API。目前已经集成入 SOK, 协同加速推荐系统 Embedding 的相关计算。
在应用了 SOK 和 HKV 后,相同规模模型和 GPU 下,陌陌精排模型的训练性能相比于原方案,整体吞吐提升了 5 倍以上。除此之外,陌陌推荐团队在当前 SOK + HKV 的架构基础上,基于业务场景特点,进一步优化整体性能,包括梯度合并,减少梯度计算的通信开销;并行特征数据读取与转换,以及特征数据预取到 GPU 等操作提速特征 IO;使用 XLA 进行编译优化,融合 kernel 以减少 kernel launch 时间;设置 GPU 亲和等操作,使得整体性能提升达到 12 倍。
在陌陌的实际应用中,动态 Embedding 的特性大大简化了推荐系统中连续训练需人工控制显存中 embedding tab 大小的问题。而 SOK 与 HKV 为陌陌提供了完整的功能和性能支持。
除此之外,我们跟陌陌的合作过程中,也结合产品部署中的需求进一步对产品进行了性能优化和功能迭代,比如:
SOK 针对 embedding vector 的拷贝和 combiner 进行了高度优化和 kernel 融合,使整个 lookup 的前后向过程拥有更好的性能。
在陌陌 GPU 高水位线的实际业务中,基于陌陌的测试和反馈,SOK 通过优化了性能和功能的平衡点,使得其在保证性能的基础上,稳定性也大大提升。
另外,在陌陌的实际应用中,面临着模型实时训练的挑战,即需要减少对计算资源的占用,因此我们引入了 HKV,它支持了层次化动态嵌入存储(包括 CPU 和 GPU),并提供了灵活的 eviction 机制以及丰富的 API。这种引入在降低资源占用的同时,也提高了系统的灵活性。
陌陌的实际业务场景和 GPU 使用方式对于 SOK 的开发和迭代提供了非常宝贵的经验,同时陌陌的大量测试反馈也帮助 SOK 提升了应对复杂场景的能力,使得 SOK 的 feature 更加稳定和贴近客户。
持续合作:优化推荐模型性能,降低训练成本
双方团队通过 SOK 和 HKV 对原方案进行深度优化后,成功帮助陌陌提升了 12 倍的训练效率,极大的降低了模型训练的成本和新模型算法尝试的成本。目前,整体方案已上线,全面支持陌陌推荐系统模型训练。
近期,NVIDIA 团队还与陌陌进行了基于 Transformer 的推荐模型性能优化。NVIDIA JOC 团队和 Merlin 团队基于客户的模型做了一系列性能分析,将 XLA+AMP+半精度 allreduce 应用到该模型上后,端到端性能实现了 50% 的加速。在此基础上,团队们进一步对性能热点 multi-head-attention 部分进行优化,正在将 Flash-Attention 以 tf-plugin 形式进行集成,预计此项优化集成后,整体加速比可达到 3 倍,同时整体的优化方案使得显存使用量下降约 70%,可以显著地缓解显存紧张的问题。
未来,陌陌与 NVIDIA 将继续在推荐系统训练和推理等方面持续合作,持续推进 GPU 和 AI 软件加速计算在陌陌的全面落地,期待能够为陌陌的业务及场景应用带来更大的价值。
了解更多本案例中相关的 NVIDIA 产品信息,敬请查阅:
- NVIDIA Merlin:
https://developer.nvidia.cn/merlin - Merlin TensorFlow Plugin (SOK) :
https://github.com/NVIDIA-Merlin/HugeCTR/tree/main/sparse_operation_kit - Merlin HierarchicalKV (HKV):
https://github.com/NVIDIA-Merlin/HierarchicalKV
GTC 2024 将于 2024 年 3 月 18 至 21 日在美国加州圣何塞会议中心举行,线上大会也将同期开放。
扫描下方海报二维码,立即注册 GTC 大会。