扩展和加速大型深度学习推荐系统 – HugeCTR 系列第 1 部分

近日，继成为计算机视觉、自然语言处理和语音识别领域的主流机器学习技术之后，深度学习又应用于推荐系统领域。
训练深度学习推荐系统并非易事，我们在[u]上一篇博文[/u]中曾探讨过这些挑战。
深度学习推荐模型通常具有一个浅层架构，几乎没有全连接层和大型用户/项目嵌入表（大小可达数 GB 或 TB）。
为应对这些挑战，NVIDIA 开发了开源框架 NVIDIA Merlin，通过 GPU 端到端（从 ETL 到训练再到部署）扩展和加速深度学习推荐系统。
这是此系列博文的第一部分，介绍了使用 HugeCTR 训练大型推荐系统的神经网络所面临的挑战以及相应的解决方案。
Merlin HugeCTR 是一个推荐系统特定框架，可加快在 GPU 上大规模训练和部署复杂的深度学习模型的速度。

嵌入表是否未超出内存？

您的第一反应可能是“是的，我想训练深度学习推荐系统，但是如何才能让嵌入表不超出 GPU 显存”。
没错，这是我们在训练大型深度学习推荐模型时遇到的一个主要挑战。
嵌入表可达 100GB 甚至数 TB，而且不适合单个 GPU 或主机内存。
HugeCTR 利用数据和模型并行性来扩展训练，并将一个嵌入表分布于多个 GPU 之上（如下图所示），从而解决了这一问题。

由于全连接层只有少数“几个”参数，因此每个 GPU 都包含一个用于数据并行训练的权重副本。
此外，还可将大型嵌入表分布于多个 GPU 或多个节点之上。
根据具体的批量，将相应地向前/向后传递，并更新所有参数。
HugeCTR 以这种方式通过模型并行（嵌入层）和数据并行（稠密层或 Dense Layer）来扩展训练。
通过实验，我们发现与 TensorFlow CPU 版本相比，速度提高了 114 倍，并且能够将嵌入表的大小扩展到 10 TB。

能否在其他深度学习框架中使用同一策略呢？

HugeCTR 已针对推荐系统进行高度优化，其核心价值在于能够在多个 GPU 上实现良好扩展。
实现这些策略并非易事，需要使用 CUDA C++ 进行自定义实现。
同时，TensorFlow 是最热门的通用深度学习框架之一，拥有强大的用户群。
我们希望向社区提供相应功能，因此将 HugeCTR 嵌入导出为 TensorFlow 自定义运算，以便它可以与客户基于 TensorFlow 构建的现有代码和管道进行互操作。
如下所示，使用 HugeCTR TensorFlow 运算只需更改一行，并且可轻松地使用 TensorFlow Keras 模型。
查看我们的 Jupyter 笔记本，了解如何通过创建新 keras 层将 HugeCTR 嵌入集成到 DeepFM 模型中。
如果您有兴趣了解更多详细信息（例如其基础机制和性能），请观看我们在 GTC（GPU 技术大会）2021 的相关演讲。
总体而言，在单个 DGX A100 上，我们的 HugeCTR TensorFlow 运算速度比原始 TensorFlow 嵌入层快了 10 倍之多。
在具有 7 个全连接层的端到端训练方面，它将速度提高了 3.5 倍。

嵌入表仍超出可用内存时可在单台设备上训练

推荐系统的嵌入表可达 100GB 甚至数 TB，即使是多 GPU 也无法满足 GPU 显存需求。
随着用户和服务项目数量的增加，表大小也会快速增加。
我们在上文讨论过，HugeCTR 可以将模型扩展为多个节点，但是如果公司的嵌入表大小超过其 GPU 总内存容量，则公司可能无法一直扩展 GPU 集群。

HugeCTR 的 Embedding Training Cache (ETC) 功能支持在单台设备上训练深度学习推荐系统，即使嵌入表超出 GPU 显存也可以。
它通过粗粒度方式解决了上述问题。
ETC 将部分冷模型存储在外存（MOS），有利于在单台设备（如 DGX A100）上训练 TB 级嵌入表。
它假设数据集被分为多个子数据集，每个子数据集都有用户指定的样本数量，并将分类特征（键集）的唯一值从每个子数据集中提取出来。

例如，用户可能想要确定各子数据集是否具有一个月的数据量，以便分类特征可以生成一个可加载给定 GPU 显存容量的嵌入表。
然后，经过多次传递完成训练，每次传递都使用不同的键集和子数据集对。

同样，model_subscriber.update() 在每次传递开始时所执行的操作类似于进程上下文切换。
它将在上一次传递中训练的嵌入向量转储到 SSD 等二级存储中。
之后，它使用当前键集将新的嵌入向量加载到设备内存，以便可以在当前传递中进行训练。
这样一来，我们就可以有效地训练大小超出可用设备内存容量的大型嵌入表了。

其理念是，将大块数据迁移到 GPU 非常高效。
对于每个子数据集，将加载所需的嵌入向量并执行多个批量更新。
之后，在参数服务器中同步模型参数。
由于嵌入向量的访问通常具有幂律分布，因此对于每个子集，我们只需要一小部分嵌入表。
在子数据集中对完整数据集进行迭代时，我们仍然会不断训练所有嵌入向量。

HugeCTR ETC 仅可在 Python 中使用。
我们提供了一个Jupyter notebook，其中描述如何基于我们的低级 Python API 启用和使用它。
我们正在积极研究如何有效地利用设备内存带宽，同时也在努力探求如何简化 SSD 和设备内存之间的路径。

试用 NVIDIA Merlin 和 HugeCTR，扩展推荐系统流水线

我们在 GitHub 存储库（NVIDIA Merlin、NVIDIA NVTabular 和 NVIDIA HugeCTR）中提供了许多适用于 HugeCTR 的示例或 ETL 训练-推理端到端示例。
有关 HugeCTR 的更多信息，请参阅我们的用户指南。
HugeCTR 提供了多种常用深度学习模型的参考实施，例如 Wide&Deep、深度交叉网络和深度学习推荐模型 (DLRM)。
如果您在使用 HugeCTR 时遇到任何问题或有任何功能请求，可以在我们的存储库中提交 GitHub 问题。
这是我们 HugeCTR 系列中关于“扩展和加速大型深度学习推荐系统”的第一篇博文。