SOK将HugeCTR中的高级特性进行包装使用

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 的的先进设计。

关于作者

NVIDIA GPU 计算专家研发团队成立于2020年7月，主要面向全球用户开发基于英伟达通用计算产品的软件解决方案。 目前团队致力于推荐系统训练框架 Merlin HugeCTR的开发。作为Merlin 系统三大核心软件之一，Merlin HugeCTR 是一套针对于推荐系统设计的模型训练框架。项目启动于2018年，目前产品应用于多家国际互联网公司，自2020年以来，HugeCTR持续参与并推动国际性竞赛MLPerf DLRM训练项目，从而实现了业界领先的优化性能。 NVIDIA GPU 计算专家研发团队是一个小型国际化工程团队，核心成员来自于多个国家（中国，韩国，荷兰，德国），客户分布于世界各地。日常工作使用英语沟通。 团队主要工作围绕软件产品开发展开，涉及产品维护、文档、测试、发布、部署，日常工作中与各国硬件、产品、工程、客户支持团队有着紧密的合作。

审核编辑：郭婷