Technical Walkthrough

使用 HugeCTR TensorFlow 嵌入插件加速嵌入

推荐系统是互联网的经济引擎。很难想象任何其他类型的应用程序会对我们的日常数字生活产生更直接的影响:数以万亿计的项目被推荐给数十亿人。推荐系统会在大量选项中过滤产品和服务,从而缓解大多数用户面临的选择悖论。

随着数据量的增加,深度学习( DL )推荐系统开始显示与传统的基于机器学习的方法相比的优势,例如梯度增强树。为了给出一个具体的数据点, NVIDIA RAPIDS 。 AI 团队与 DL 赢得了三场推荐比赛最近:

甚至在一年前 NVIDIA 数据科学家询问为什么深度学习模型还没有在推荐系统竞赛中持续获胜?时,这种情况也没有持续发生。

嵌入在现代基于 DL 的推荐体系结构中起着关键作用,为数十亿实体(用户、产品及其特征)编码单个信息。随着数据量的增加,嵌入表的大小也随之增加,现在这些表跨越多个 GB 到 TB 。在训练这种类型的 DL 系统时存在着独特的挑战,其庞大的嵌入表和稀疏的访问模式可能跨越多个 GPU 节点(如果不是节点的话)。

本文重点介绍 NVIDIA Merlin 推荐系统框架如何应对这些挑战,并介绍了一种优化的嵌入实现,其性能比其他框架的嵌入层高出 8 倍。这个优化的实现也可以作为一个 TensorFlow 插件提供,它可以与 TensorFlow 无缝地工作,并作为 TensorFlow 本机嵌入层的方便替代品。

Embeddings

Embeddings 是一种机器学习技术,它将每个感兴趣的对象(用户、产品、类别等)表示为密集的数字向量。因此,嵌入表只是一种特定类型的键值存储,键值是用于唯一标识对象的 ID ,值是实数向量。

Embeddings 是现代 DL 推荐系统中的一个关键构建块,通常位于输入层之后、“特征交互”和密集层之前。嵌入层是从数据和端到端训练中学习的,就像深层神经网络的其他层一样。正是嵌入层将 DL recommender 模型与其他类型的 DL 工作负载区分开来:它们为模型提供了大量的参数,但几乎不需要计算,而计算密集型密集层的参数数量要少得多。

举一个具体的例子:原始博大精深模型有几个密集的层,大小为[1024 、 512 、 256],因此只有几百万个参数,而其嵌入层可以有数十亿个条目和数十亿个参数。例如,这与 NLP 领域流行的 BERT 模型体系结构形成对比,其中嵌入层只有数万个条目,总计数百万个参数,但密集的前馈和注意层由数亿个参数组成。这种差异还导致了另一个观察:与其他类型的 DL 模型相比, DL recommender 网络每字节输入数据的计算量通常要小得多。

为什么优化嵌入对推荐者工作流很重要

为了理解为什么优化嵌入层和相关操作很重要,以下是培训嵌入的挑战:大小和访问速度。

Size

随着在线平台和服务获得数亿甚至数十亿用户,以及提供的独特产品数量达到数十亿,嵌入表的规模不断扩大也就不足为奇了。

据报道 Instagram 已被正在开发大小达到 10 TB 的推荐机型删除。同样,百度报告了广告排名模型也达到了 10TB 的境界。在整个行业中,数百 GB 到 TB 的型号越来越流行,例如Pinterest 的 4-TB 模型谷歌的 1 . 2 TB 模型

很自然,在单个计算节点上拟合 TB 规模的模型是一个重大挑战,更不用说在单个计算加速器上了。作为参考,最大的 NVIDIA A100 GPU 目前配备了 80 GB 的 HBM 。

访问速度

训练推荐系统本质上是一项内存带宽密集型任务。这是因为每个训练样本或批通常在嵌入表中涉及少量实体。必须检索这些条目以计算向前传递,然后在向后传递中更新。

CPU 主存储器具有高容量,但带宽有限,高端型号通常在几十 GB / s 的范围内。另一方面, GPU 的内存容量有限,但带宽很高。 NVIDIA A100 80-GB GPU 提供 2 TB / s 的内存带宽。

解决方案

这些挑战以不同的方式得到解决。例如,将整个嵌入表保留在主存上可以解决大小问题。然而,它通常会导致极低的训练吞吐量,这往往与新数据的数量和速度相形见绌,从而使系统无法及时重新训练。

或者,嵌入可以被仔细地分散在多个 GPU S 和多个节点上,仅被通信瓶颈所困扰,导致持续严重的 GPU – 在使用和训练性能下的计算与纯 CPU 训练正好一致。

嵌入层是推荐系统的主要瓶颈之一。优化嵌入层是解锁 GPU 高计算吞吐量的关键。

在下一节中,我们将讨论 NVIDIA Merlin HugeCTR 推荐框架如何通过使用 NVIDIA 技术,如 GPU 直接远程直接内存访问( RDMA )、 NVIDIA 集体通信图书馆( NCCL )、 NVLink 和 NVSwitch ,解决大规模嵌入式的挑战。它解锁了 GPU 的高计算和高带宽容量,同时解决了开箱即用、多 GPU 、多节点支持和模型并行性的内存容量问题。

NVIDIA Merlin HugeCTR 嵌入概述

NVIDIA Merlin 解决了培训大规模推荐系统的挑战。它是一个端到端的推荐框架,可以加速推荐系统开发的所有阶段,从数据预处理到培训和推理。 NVIDIA Merlin HugeCTR 是一个开源的推荐系统,专用 DL 框架。在这篇文章中,我们关注 HugeCTR 的一个特定方面:嵌入优化。

有两种方法可以利用 HugeCTR 中的嵌入优化工作:

  • 将本机 NVIDIA Merlin HugeCTR 框架用于培训和推理工作负载
  • 使用 NVIDIA Merlin HugeCTR TensorFlow 插件,该插件旨在与 TensorFlow 无缝协作

原生 HugeCTR 嵌入优化

为了克服嵌入挑战并实现更快的训练, HugeCTR 实现了自己的嵌入层,其中包括 GPU 加速哈希表、以节省内存的方式实现的高效稀疏优化器以及各种嵌入分布策略。它利用NCCL作为其内部 GPU 通信原语。

哈希表实现基于 RAPIDS cuDF ,它是 GPU 数据帧库,构成 NVIDIA 的 RAPIDS 数据科学平台的一部分。 cuDF GPU 哈希表可以实现比 CPU 实现高达 35 倍的加速,例如来自线程构造块( TBB )的并发_ hash _映射。有关更多信息,请参阅介绍 NVIDIA Merlin HugeCTR :一个专门用于推荐系统的培训框架

考虑到可伸缩性, HugeCTR 默认支持嵌入层的模型并行性。嵌入表分布在可用的 GPU 和节点上。另一方面,密集层采用数据并行(图 1 )。

HugeCTR employs model parallelism for the embedding layer by default. The embedding tables will be distributed across the available GPUs
图 1 。 HugeCTR 中的嵌入层并行性

腾讯推荐团队是原生 HugeCTR 框架的首批采用者之一,大量使用了原生嵌入层。腾讯广告与深度学习平台负责人孔祥亭在最近的一次interview中表示,“ HugeCTR 作为一个推荐培训框架,被整合到[腾讯]广告推荐培训系统中,以加快模型培训的更新频率,并可以培训更多样本以提高在线效果。”

HugeCTR TensorFlow 插件

NVIDIA Merlin 框架的所有组件都是开源的,旨在与更大的深度学习和数据科学生态系统进行互操作。我们的长期愿景是加速 GPU 上的推荐工作负载,而不管您喜欢哪种框架。 HugeCTR TensorFlow 嵌入插件是为了实现这一目标而创建的。


在高层次上, TensorFlow 嵌入插件是通过利用许多与本机 HugeCTR 嵌入层相同的嵌入优化技术来设计的。特别是,这将是 GPU 哈希表和 NCCL ,位于引擎盖下,用于 GPU 之间的通信。

HugeCTR 嵌入插件设计为方便无缝地与 TensorFlow 配合使用,作为 TensorFlow 本机嵌入层(如tf.nn.embedding_lookuptf.nn.embedding_lookup_sparse)的替代品。它还提供了开箱即用的高级功能,例如在多个 GPU 上分布嵌入表的模型并行性。

NVIDIA Merlin HugeCTR TensorFlow 插件演练

下面是如何使用 TensorFlow 嵌入插件。完整示例可在 HugeCTR repository上找到,我们还提供了完整的基准测试笔记本用于再现性能数据。

访问 HugeCTR 嵌入插件最方便的方式是使用 NGC NVIDIA Merlin TensorFlow 培训码头工人形象,在该插件中预编译和安装 NVIDIA Merlin 框架的其他组件以及 TensorFlow 。最新的版本可以直接从 HugeCTR 存储库中提取,实时编译和安装。更新 TensorFlow 时,还必须为新安装的 TensorFlow 版本重新编译插件。

为了进行比较,这里介绍了如何使用本机 TensorFlow 嵌入层。首先,初始化 2D 数组变量以保存嵌入的值。然后,使用 tf . nn . embedding _ lookup 查找 ID 列表对应的嵌入值。

embedding_var = tf.Variable(initial_value=initial_value, dtype=tf.float32, name='embedding_variables')

@tf.function
def _train_step(inputs, labels):
    emb_vectors = tf.nn.embedding_lookup([self.embedding_var], inputs)
    ...

for i, (inputs, labels) in enumerate(dataset):
    _train_step(inputs)

同样,也可以使用 HugeCTR 嵌入插件。首先,初始化嵌入层。接下来,该嵌入层用于查找 ID 列表的相应嵌入值。

import sparse_operation_kit as sok

emb_layer = sok.All2AllDenseEmbedding(max_vocabulary_size_per_gpu,
                                      embedding_vec_size,
                                      slot_num, nnz_per_slot)

@tf.function
def _train_step(inputs, labels):
    emb_vectors = emb_layer(inputs)
    ...

for i, (inputs, labels) in enumerate(dataset):
    _train_step(inputs)

HugeCTR 嵌入插件旨在与 TensorFlow 无缝协作,包括 Adam 和 sgd 等其他层和优化器。在 TensorFlow v2 . 5 之前, Adam 优化器是基于 CPU 的实现。

为了充分发挥 HugeCTR 嵌入插件的潜力,我们还在 sok . optimizers . adam 中提供了一个基于 GPU 的插件版本。从 TensorFlow v2 . 5 开始,标准 Adam 优化器 tf . keras . optimizers . Adam (现在带有 GPU 实现)可以以类似的精度和性能使用。

绩效基准

在本节中,我们通过合成和实际用例展示 HugeCTR TensorFlow 嵌入插件的性能。

合成数据

在本例中,我们使用具有 100 个特征字段的合成数据集,每个都有 10 个查找,词汇表大小为 8192 。推荐模型是一个具有六层的 MLP ,每层大小为 1024 。使用 TensorFlow 中的精确模型架构、优化器和数据加载器,我们观察到在 1x A100 GPU 上, HugeCTR 嵌入插件比本机 TensorFlow 嵌入层的平均迭代时间提高了 7 . 9 倍(图 2 )。

当从 1 到 4 个 A100 HugeCTR s 进行强缩放时,我们观察到总加速比为 23 . 6 倍。默认情况下, GPU 嵌入插件提供了多重 GPU 缩放的好处。在引擎盖下,嵌入插件以模型并行方式自动将与功能字段对应的表格分发到可用的 GPU 上。这与本机 TensorFlow 嵌入层形成对比,在本机嵌入层中,分布式模型并行多 GPU 训练需要大量额外工作。 TensorFlow 分布策略 Mirrored 策略和 MultiWorkerMirrored 策略都是为进行数据并行同步训练而设计的。

HugeCTR TensorFlow plugin provides a 7.9x speedup over native TensorFlow 2.5 embedding lookup layer.
图 2 . HugeCTR TensorFlow 在合成数据集上 NVIDIA DGX A100 80GB 上嵌入插件的性能

真实用例:美团推荐系统

美团推荐系统团队是首批成功采用 HugeCTR TensorFlow 插件的团队之一。起初,团队基于 CPU 优化了他们的培训框架,但随着他们的模型变得越来越复杂,很难对培训框架进行更深入的优化。现在,美团正致力于将 NVIDIA HugeCTR 集成到基于 A100 GPU 的培训系统中。

美团高级技术专家黄军( Jun Huang )表示:“一台配备 8x A100 GPU s 的服务器可以替代基于 CPU 的培训系统中的数百名员工。成本也大大降低。这是初步的优化结果,未来仍有很大的优化空间。”。

美团使用DIEN作为推荐模型。嵌入参数的总数是数百亿,每个样本中有数千个特征字段。由于输入特征的范围事先不是固定的和未知的,团队使用哈希表在输入嵌入层之前唯一地标识每个输入特征。

使用 TensorFlow 中的精确模型架构、优化器和数据加载器,我们观察到在单个 A100 GPU 上, HugeCTR 嵌入插件实现了比原始 TensorFlow 嵌入 11 . 5 倍的加速。在弱缩放情况下, 8x A100 GPU s 上的迭代时间仅略微增加到 1x A100 GPU 的 1 . 17 倍(图 3 )。

On a real use case, HugeCTR TensorFlow plugin provides a 11.6x speedup over native TensorFlow 2.5 embedding lookup layer.
图 3 .美团数据上 NVIDIA DGX A100 80GB 上的 HugeCTR TensorFlow 嵌入插件性能

结论

HugeCTR TensorFlow 嵌入插件今天可以从 HugeCTR GitHub 存储库以及 NGC NVIDIA Merlin TensorFlow 容器获得。如果您是 TensorFlow 用户,希望构建和部署具有大型嵌入表的大规模推荐系统, HugeCTR TensorFlow 插件可以轻松替换 TensorFlow 嵌入查找层。

尝试一下,看看 GPU 的全部潜力。如果您觉得需要更高的性能和优化,那么您下一步要尝试的或许是成熟的原生 HugeCTR 框架。