摘要
1. Introduction
HugeCTR 能够高效地利用 GPU 来进行推荐系统的训练,为了使它还能直接被其他 DL 用户,比如 TensorFlow 所直接使用,我们开发了 SparseOperationKit (SOK),来将 HugeCTR 中的高级特性封装为 TensorFlow 可直接调用的形式,从而帮助用户在 TensorFlow 中直接使用 HugeCTR 中的高级特性来加速他们的推荐系统。
图 1. SOK embedding 工作流程
SOK 以数据并行的方式接收输入数据,然后在 SOK 内部做黑盒式地模型转换,最后将计算结果以数据并行的方式传递给初始 GPU。这种方式可以尽可能少地修改用户已有的代码,以更方便、快捷地在多个 GPU 上进行扩展。
SOK 不仅仅是加速了 TensorFlow 中的算子,而是根据业界中的实际需求提供了对应的新解决方案,比如说 GPU HashTable。SOK 可以与 TensorFlow 1.15 和 TensorFlow 2.x 兼容使用;既可以使用 TensorFlow 自带的通信工具,也可以使用 Horovod 等第三方插件来作为 embedding parameters 以为参数的通信工具。
2. TF2 Comparison/Performance
使用 MLPerf 的标准模型 DLRM 来对 SOK 的性能进行测试。
图 2. SOK性能测试数据
相比于 NVIDIA 的 DeepLearning Examples,使用 SOK 可以获得更快的训练速度以及更高的吞吐量。
3. API
SOK 提供了简洁的、类 TensorFlow 的 API;使用 SOK 的方式非常简单、直接;让用户通过修改几行代码就可以使用 SOK。
- 定义模型结构
左侧是使用 TensorFlow 的 API 来搭建模型,右侧是使用 SOK 的 API 来搭建相同的模型。使用 SOK 来搭建模型的时候,只需要将 TensorFlow 中的 Embedding Layer 替换为 SOK 对应的 API 即可。
- 使用 Horovod 来定义 training loop
同样的,左侧是使用 TensorFlow 来定义 training loop,右侧是使用 SOK 时,training loop 的定义方式。可以看到,使用 SOK 时,只需要对 Embedding Variables 和 Dense Variables 进行分别处理即可。其中,Embedding Variables 部分由 SOK 管理,Dense Variables 由 TensorFlow 管理。
- 使用 tf.distribute.MirroredStrategy 来定义 training loop
类似的,还可以使用 TensorFlow 自带的通信工具来定义 training loop。
- 开始训练
在开始训练过程时,使用 SOK 与使用 TensorFlow 时所用代码完全一致。
4. 结语
SOK 将 HugeCTR 中的高级特性包装为 TensorFlow 可以直接使用的模块,通过修改少数几行代码即可在已有模型代码中利用上 HugeCTR 的的先进设计。
更多信息,请参考 SOK 官方文档:
以下是 HugeCTR 的 Github repo 以及其他发布的文章,欢迎感兴趣的朋友阅读和反馈。
Github: https://github.com/NVIDIA-Merlin/HugeCTR (更多文章详见 README)
除此之外,NVIDIA Merlin HugeCTR 团队正在积极招募 C++ 以及 CUDA 工程师,欢迎各位有意向的同学发邮件至 jershi@nvidia.com 或 yincanw@nvidia.com 积极申请!JD 请点击:NVIDIA GPU 计算专家研发团队招新啦!你将成为 Merlin 应用框架的核心成员!查看详细介绍!。