[源码解析] NVIDIA HugeCTR，GPU 版本参数服务器 --(1)

[源码解析] NVIDIA HugeCTR，GPU版本参数服务器 --(1)

0x00 摘要

本系列之中我们将会介绍 NVIDIA 出品的 HugeCTR，这是一个面向行业的推荐系统训练框架，针对具有模型并行嵌入和数据并行密集网络的大规模 CTR 模型进行了优化。

本文以Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems，GitHub 源码文档 https://github.com/NVIDIA-Merlin/HugeCTR 的翻译为基础，并且结合源码进行分析。

其中借鉴了HugeCTR源码阅读这篇大作，特此感谢，期望能在此篇大作基础之上，再丰富一下对HugeCTR的理解。

0x01 背景

我们将简要讨论 CTR 估计在现代推荐系统中的作用及其训练中的主要挑战。

1.1 推荐系统中的点击率估计

从在线广告和电子商务到流媒体服务，推荐系统无处不在，同时对服务提供商的收入产生巨大影响。推荐系统找到给定用户最可点击的项目，然后对它们进行排名并向用户显示前 N 个项目。为了实现这个目标，推荐系统首先必须估计特定用户点击项目的可能性。此任务通常称为 CTR 估计。

如何估算点击率？这里没有巫术，一般是获取包含用户-物品交互的富数据集，并使用它来训练 ML 模型。数据集中的每条记录都可以包含来自用户（年龄、工作），商品（类型、价格）和用户商品点击（0 或 1）的特征。例如，如果用户 A 从一系列书籍中购买或点击了几本传记，那么模型为传记分配高概率值是有意义的。

CTR 的系统结构大致如下：

下图展示了CTR推理流程。

图来自HugeCTR_Webinar

1.2 点击率估算训练的挑战

首先，推荐系统之中的特征有如下性质：高维，稀疏。大规模推荐系统会面临用户和物品的频繁变化，因此识别用户点击背后的隐式特征交互至关重要，这样推荐系统可以提供更高质量的更通用的推荐。例如，30 岁以下的已婚人士和孩子未满 2 岁的人可能倾向于购买高 ABV 的啤酒。对这些隐式特征交互进行建模需要领域专家进行复杂的特征工程。更糟糕的是，由于特征极其复杂且不直观，即使是人类专家也常常无法发现这些交互。为了代替这种对专家的依赖，人们研究出了一些基于深度学习的方法，例如 Wide & Deep，DeepFM 和 DLRM，这些模型可以捕获这些复杂的交互。

训练 CTR 估计模型的另一个挑战是用户和物品几乎每天都在变化，因此训练出来的模型其生命周期可能很短。此外，由于数据集的大小的增加，维数和稀疏性因素，CTR 模型通常包含一个很大的嵌入表，其可能无法放入单个 GPU 甚至多个 GPU 的节点中。因此，数据加载，嵌入表查找和 GPU 间通信可以占据模型训练时间的很大一部分。

这些因素，再加上缺乏用于 CTR 估算的标准化建模方法，通常导致服务在吞吐量和延迟方面经常只能达到次优性能。所以在单个或多个 GPU 上完成模型的更快迭代训练是非常重要的。

0x02 HugeCtr

HugeCTR 是一个开源框架，用于在 NVIDIA GPU 上加速 CTR 估计模型的训练，并针对 NVIDIA GPU 的性能进行了高度优化，同时允许用户以 JSON 格式自定义模型。它是用 CUDA C++ 编写的，并且高度利用了 GPU 加速库，例如cuBLAS、cuDNN和NCCL。它最初是作为内部原型来评估 GPU 在 CTR 估计问题上的潜力，但是其很快成为基于 GPU 的推荐系统的参考设计。由于它自然而然地成为了专用于 CTR 估算的更通用的框架，因此 NVIDIA 于 2019 年 9 月开源了其初始版本，以接受外部反馈，同时与一些客户保持互动。

HugeCTR 也是 NVIDIA Merlin的支柱，这是一个框架和生态系统，用于构建需要大量数据集进行训练的大规模推荐系统，旨在促进推荐系统开发的所有阶段，并在 NVIDIA GPU 上加速。

图来自源码 https://github.com/NVIDIA-Merlin/Merlin

HugeCTR 在单个 NVIDIA V100 GPU 上的速度比 TensorFlow 在 40 核 CPU 节点上提高了 114 倍，在同一个 V100 GPU 上实现了 TensorFlow 的 8.3 倍提高。由于由线性模型和深度模型组成的混合模型已变得普遍，因此 HugeCTR 架构 2.1 版扩展为支持 Wide & Deep、DCN 和 DeepFM 等模型。更新包括新的数据读取器，它可以同时读取连续和分类输入数据；以及新的层，包括因子分解机和交叉层。为了实现更灵活的设计空间探索，还添加了 Dropout、L1/L2 正则化器等。

0x03 架构

3.1 CTR DL 模型

下图描绘了用于 CTR 估计的 DL 模型的步骤：

按批次读取数据记录，每个记录都由高维、极其稀疏（或 categorical 类型）的特征组成。每个记录还可以包含密集的数字特征，这些特征可以直接馈送到全连接层。
使用嵌入层将输入稀疏特征压缩为低维密集嵌入向量。例如，如果有 N 个稀疏特征，嵌入维度为 K，则嵌入表生成 N 个 K 维密集向量。
使用前馈神经网络来估计点击率。

图上显示了一个典型的 CTR 模型，包括数据读取器、嵌入和全连接层。图来自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

3.2 HugeCTR 架构

HugeCTR 不仅支持 CTR DL 所有三个步骤，而且还增强了端到端的性能，比如：

为了防止数据加载成为训练中的主要瓶颈，它实现了一个专用的数据读取器，该读取器是异步和多线程的。它将读取一组批处理数据记录，其中每条记录都由高维、极度稀疏或分类特征（categorical features）组成。每个记录还可以包含密集的数字特征（dense numerical features），这些特征可以直接馈送到全连接层。
嵌入层用于将稀疏输入特征压缩为低维、密集的嵌入向量。共有三个 GPU 加速的嵌入阶段：
- 表查找
- 每个插槽（slot）内的权重规约。
- 跨插槽的权重拼接（concatenation）。
通过利用高效的 CUDA 优化技术和支持 CUDA 的库来支持前向和后向传播中的所有层，优化器和损失函数都是在 CUDA C++ 中实现的。

为了训练大规模 CTR 估计模型，HugeCTR 中的嵌入表是模型并行的，并分布在同构集群中的所有 GPU 上，该集群由多个节点组成。每个 GPU 都有自己的：

前馈神经网络（数据并行）来估计点击率。
哈希表使数据预处理更容易并启用动态插入。

所以，可以扩展到多个 GPU 和节点的HugtCTR的架构总结如下：

3.3 基于GPU的参数服务器

HugeCTR 实现的是一个基于GPU的参数服务器，其将embedding层放到GPU之中，worker通过与参数服务器的交互来获取embedding。

图来自HugeCTR_Webinar

0x04 核心功能

在本节中，我们将介绍 HugeCTR 的关键特性，这些特性有助于其高性能和可用性。注意：多节点训练和混合精度训练可以同时使用。

4.1 模型并行训练

HugeCTR 原生支持模型并行和数据并行训练，使得在 GPU 上训练非常大的模型成为可能。

4.1.1 in-memory GPU hash table

在 CTR 估计中，嵌入（embedding）对于获得不错的模型精度几乎是必不可少的。它通常会导致对内存容量和带宽的高需求以及相当数量的并行性。如果embedding分布在多个 GPU 或多个节点上，则通信开销也可能很大。由于用户和物品数量庞大且不断增加，庞大的嵌入表在所难免。

为了克服这些挑战并实现更快的训练，HugeCTR实现了自己的嵌入层，其中包括一个 GPU 加速的哈希表，并利用NCCL 作为其 GPU 间通信原语。哈希表的实现基于RAPIDS cuDF 的实现，RAPIDS cuDF 是来自 NVIDIA 的 GPU DataFrame 库。cuDF GPU 哈希表可以比 Threading Building Blocks (TBB) 的 concurrent_hash_map 多出高达 35 倍的加速。

总之，HugeCTR 支持跨越同构计算集群中的多个 GPU 和多个节点的模型并行嵌入表。嵌入的特征和类别可以分布在多个 GPU 和节点上。例如，如果您有两个具有 8xA100 80GB GPU 的节点，则可以完全在 GPU 上训练大至 1TB 的模型。通过使用嵌入训练缓存，您可以在相同节点上训练更大的模型。

4.1.2 Multi-slot embedding

嵌入表可以被分割成多个槽（或feature fields）。在嵌入查找过程中，属于同一槽的稀疏特征输入在分别转换为相应的密集嵌入向量后，被简化为单个嵌入向量。然后，来自不同槽的嵌入向量连接在一起。

多槽（multi-slot）嵌入通过以下方式提高了 GPU 间带宽利用率：

当数据集中有很多特征时，它有助于将每个槽中有效特征的数量减少到可管理的程度。
通过拼接不同插槽的输出，它减少了 GPU 之间的事务数量，从而促进了更高效的通信。

下图显示了操作序列和 GPU 间通信 ( all2all) 是如何发生的。

该图显示了一个跨越 4 个 GPU 的模型并行嵌入，以及它如何与这些 GPU 的神经网络进行交互。它还显示了如何减少每个插槽的输入特征并跨两个插槽连接

该图显示了一个跨越 4 个 GPU 的模型并行嵌入，以及它如何与这些 GPU 的神经网络进行交互。它还显示了如何减少每个插槽的输入特征并跨两个插槽连接。图来自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

多槽嵌入对线性模型也很有用，它基本上是特征的加权和，只需将槽数和嵌入维度都设置为 1 即可。有关更多信息，请参阅Wide & Deep 示例。

4.1.3 具体实现

为了在不同的嵌入上获得最佳性能，可以选择不同的嵌入层实现。这些实现中的每一个都针对不同的实际培训案例，例如：

LocalizedSlotEmbeddingHash：同一个槽（特征域）中的特征会存储在一个GPU中，这就是为什么它被称为“本地化槽”，根据槽的索引号，不同的槽可能存储在不同的GPU中。LocalizedSlotEmbedding 针对每个embedding 小于 GPU 内存大小的实例进行了优化。由于在 LocalizedSlotEmbedding 中使用了每个插槽的局部规约（查完 embedding 得到向量之后，因为已经拿到了这个slot 的所有 embedding，可以做完pooling之后再做多GPU卡通信），而在 GPU 之间没有全局规约，因此 LocalizedSlotEmbedding 中的整体数据传输量远小于 DistributedSlotEmbedding。

注意：确保输入数据集中没有任何重复的键。
DistributedSlotEmbeddingHash：所有特征都存储于不同特征域/槽上，不管槽索引号是多少，这些特征都根据特征的索引号分布到不同的GPU上。这意味着同一插槽中的特征可能存储在不同的 GPU 中，这就是将其称为“分布式插槽”的原因。由于需要全局规约，所以 DistributedSlotEmbedding 适合 embedding 大于 GPU 内存大小的情况，因而 DistributedSlotEmbedding 在 GPU 之间有更多的内存交换。

注意：确保输入数据集中没有任何重复的键。
LocalizedSlotEmbeddingOneHot：一种特殊的 LocalizedSlotEmbedding，需要一个独热数据输入。每个特征字段也必须从零开始索引。例如，性别应该是0,1，而1,2 就不正确。

一定要注意，LocalizedSlotEmbeddingHash 和 DistributedSlotEmbeddingHash 的区别在于同一个槽（特征域）中的特征 是不是 会存储在同一个GPU中。比如，有 2 张GPU卡，有4个slot。

local 模式：GPU0 存 slot0 和 slot1，GPU1 存 slot2 和 slot3。
distribute 模式：每个 GPU 都会存所有 slot 的一部分参数，通过哈希方法决定如何将一个参数分配到哪个 GPU 上。

4.2 多节点训练

多节点训练使得我们很容易训练任意大小的嵌入表。在多节点解决方案中，稀疏模型（称为嵌入层）分布在节点之间。同时，密集模型（例如 DNN）是数据并行的，并且在每个 GPU 中都包含密集模型的副本（见下图）。通过我们的实施，HugeCTR 利用 NCCL 进行高速和可扩展的节点间和节点内通信。

图来自源码。

要在多个节点上运行，HugeCTR 应该使用 OpenMPI 构建。建议支持GPUDirect RDMA以获得高性能。有关更多信息，请参阅DCN 多节点训练样本。

4.3 混合精度训练

混合精度训练已成为在保持模型精度的同时实现进一步加速的常用技术，可以帮助我们改善和减少内存吞吐量占用。在 HugeCTR 中，可以配置全连接层以利用 NVIDIA Volta 架构及其后续架构上的张量核心。它们在内部使用 FP16 进行加速矩阵乘法，但其输入和输出仍为 FP32。

混合精度训练在这种模式下，TensorCores 被用于提高基于矩阵乘法的层的性能，例如FullyConnectedLayer和InteractionLayer，在 Volta、Turing 和 Ampere 架构上。对于包括嵌入在内的其他层，数据类型更改为 FP16，以便节省内存带宽和容量。要启用混合精度模式，请在配置文件中指定 mix_precision 选项。当mixed_precision设定，完整的FP16管道将被触发。将应用损失缩放以避免算术下溢（见图）。可以使用配置文件启用混合精度训练。

图 5：算术下溢图来自源码。

4.4 SGD 优化器和学习率调度

学习率调度允许用户配置其超参数，包括以下内容：

learning_rate：基础学习率。
warmup_steps：用于预热的初始步骤数。
decay_start：指定学习率衰减开始的时间。
decay_steps：衰减期（逐步）。

图 6 说明了这些超参数如何与实际学习率相互作用。

有关更多信息，请参阅Python 接口。

图 6：学习率调度图来自源码。

4.5 嵌入训练缓存

嵌入训练缓存（Model Oversubscription）使您能够训练高达 TB 的大型模型。它是通过在训练阶段以粗粒度、按需方式将超过 GPU 内存聚合容量的嵌入表的一个子集加载到 GPU 中来实现的。要使用此功能，您需要将数据集拆分为多个子数据集，同时从中提取唯一键集（见图 7）。

此功能目前支持单节点和多节点训练。它支持所有嵌入类型，并且可以与Norm和Raw数据集格式一起使用。我们修改了我们的criteo2hugectr工具以支持 Criteo 数据集的密钥集提取。有关更多信息，请参阅我们的Python Jupyter Notebook，了解如何将此功能与 Criteo 数据集结合使用。

注意：Criteo 数据集是一个常见用例，但模型预取不限于此数据集。

Fig. 7: Preprocessing of dataset for model oversubscription 图来自源码。

4.6 HugeCTR 到 ONNX 转换器

HugeCTR to Open Neural Network Exchange (ONNX) 转换器是一个hugectr2onnxPython 包，可以将 HugeCTR 模型转换为 ONNX。它可以提高 HugeCTR 与其他深度学习框架的兼容性，因为 ONNX 作为 AI 模型的开源格式。

使用我们的 HugeCTR Python API 进行训练后，您可以获得密集模型、稀疏模型和图形配置的文件，这些文件在使用该hugectr2onnx.converter.convert方法时需要作为输入。每个 HugeCTR 层将对应一个或多个 ONNX 算子，训练好的模型权重将作为初始化器加载到 ONNX 图中。此外，您可以选择使用convert_embedding标志转换稀疏嵌入层。

4.7 分层参数服务器

HugeCTR 分层参数服务器 (POC) 上的本地 SSD 和 CPU 内存之间实现了分层存储机制。通过这种实现，嵌入表不再需要存储在本地 CPU 内存中。添加了分布式 Redis 集群作为 CPU 缓存，以存储更大的嵌入表并直接与 GPU 嵌入缓存交互。为了帮助 Redis 集群查找丢失的嵌入键，已实现本地 RocksDB 作为查询引擎来备份本地 SSD 上的完整嵌入表。

4.8 异步多线程数据管道

如果没有高效的数据管道，即使向前和向后传播以光速运行，其效果也如同到达机场的时间远长于飞行时间。另外，当数据集很大并且经常变化时，将其拆分为多个文件是非常合理的。

为了有效地把数据获取这个长延迟隐藏起来，HugeCTR 有一个多线程数据读取器，其可以将数据获取与实际模型训练重叠起来。如下图所示，DataReader是一个façade，由多个并行工作器和一个收集器组成。

每个工作器每次从其分配到的数据集文件中读取一个批次。收集器会将收集到的数据记录分发到多个 GPU。所有的工作人员、收集器和模型训练作为不同的线程在 CPU 上同时运行。

该图显示了四个数据读取器如何将数据从磁盘读取到主机内存，一个收集器如何读取其中之一以提供给模型训练管道。

Figure 4. HugeCTR multithreaded data reader.

图来自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

下图显示了 HugeCTR 流水线如何把 "数据从磁盘读取到 CPU 内存的数据"，"从 CPU 到 GPU 的数据传输"以及"在 GPU 上跨不同批次的实际训练"这三个阶段重叠起来。

此图显示了“读取文件”、“复制到 GPU”和“训练”阶段如何重叠三个批次以提高 GPU 资源利用率。

图来自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems

4.9 灵活模型配置

尽管 CTR 模型之间存在一些共性，但它们的细节（包括超参数）可能有所不同。为了实现模型的灵活定制，HugeCTR 允许以 JSON 格式直观地配置模型。

例如，要描述如下图所示的混合模型，您可以编写如图 (b) 中抽象所示的“layers”子句。您可以有多个嵌入，您还可以指定批处理大小、优化器、数据路径等。在同一个配置文件中，您也可以指定用于训练的 GPU 数量和数量。有关更多信息，请参阅HugeCTR 用户指南和示例配置文件。

Figure 6. A hybrid model with two embeddings and two different types of inputs. (a) An example mode expressible by HugeCTR. (b) The corresponding config. A lot of details are omitted for simplicity.

图来自Introducing NVIDIA Merlin HugeCTR: A Training Framework Dedicated to Recommender Systems