0
  • 聊天消息
  • 系统消息
  • 评论与回复
登录后你可以
  • 下载海量资料
  • 学习在线课程
  • 观看技术视频
  • 写文章/发帖/加入社区
会员中心
创作中心

完善资料让更多小伙伴认识你,还能领取20积分哦,立即完善>

3天内不再提示

如何在TensorFlow2中高效培训1130亿参数推荐系统

星星科技指导员 来源:NVIDIA 作者:Tomasz Grel 2022-04-19 15:12 次阅读

深度学习推荐系统通常使用大型嵌入表。很难将它们放入 GPU 内存中。

这篇文章向你展示了如何结合使用模型并行和数据并行训练范例来解决这个记忆问题,从而更快地训练大型深度学习推荐系统。我分享了我的团队在 TensorFlow 2 中高效培训 1130 亿参数推荐系统所采取的步骤,该模型的所有嵌入的总大小为 421 GiB 。

通过在 GPU 和 CPU 之间拆分模型和嵌入,我的团队实现了 43 倍的加速。然而,将嵌入分布到多个 GPU 上,带来了令人难以置信的 672 倍的加速。这种多 GPU 方法实现了显著的加速,使您能够在几分钟内而不是几天内训练大型推荐系统。

您可以使用 NVIDIA 深度学习示例 GitHub 存储库 中提供的代码自己复制这些结果。

嵌入层的模型并行训练

在数据并行训练中,每个 GPU 存储模型的相同副本,但在不同的数据上训练。这对于许多深度学习应用程序来说都很方便,因为它易于实现,并且通信开销相对较低。然而,这种模式要求神经网络的权重适合单个设备。

如果模型大小大于设备内存,一种方法是将模型分成子部分,并在不同的 GPU 上训练每个子部分。这被称为模型并行训练。

表的每一行对应于要映射到密集表示的输入变量的值。表中的每一列表示输出空间的不同维度,表示所有向量中一个值的切片。因为一个典型的深度学习推荐程序会吸收多个分类特征,所以它需要多个嵌入表。

对于具有多个大型嵌入的推荐程序,有三种实现模型并行性的方法:

Table-wise split——每个嵌入表完全放在一个设备上;每个设备只包含所有嵌入的一个子集。(图 1 )

Column-wise split–每个 GPU 包含每个嵌入表中的一个子集列。(图 2 )

Row-wise split–每个 GPU 保存每个嵌入表中的行子集。

由于负载平衡问题,行分割比其他两个选项更难实现。在本文中,我将重点介绍表拆分和列拆分。混合和匹配多种方法是一个可行的选择,但为了简单起见,我不会在本文中集中讨论这一点。

pYYBAGJeYNOAN3PvAAEgyPGTFXU523.png

图 1 。按表拆分模式是指每个 GPU 存储所有嵌入表的子集

图 2 。按列拆分模式是指每个设备存储来自每个嵌入表的列的子集

这些方法之间有一些关键区别(表 1 )。简言之,按表拆分模式更易于使用,而且可能更快,具体取决于具体的工作负载。

一个缺点是它不支持嵌入跨越多个 GPU 的表。相比之下,按列拆分模式支持嵌入跨多个 GPU 的表,但速度可能会稍慢,尤其是对于窄表。

表 1 。表拆分和列拆分模式之间的比较。

高效训练推荐系统的混合并行方法

典型的推荐程序在嵌入后运行算术密集型层,如线性或点积。处理模型这一部分的一种幼稚方法是将嵌入查找的结果收集到单个 GPU 上,并在此 GPU 上运行这些密集层。然而,这是非常低效的,因为在这段时间内没有使用用于保存嵌入的另一个 GPU 。

更好的方法是使用所有 GPU 通过数据并行运行密集层。这可以通过按批量大小拆分嵌入查找的结果来实现。也就是说,对于 N 和八 GPU 的全局批量,每个 GPU 只处理 N / 8 个训练样本。实际上,这意味着密集层以数据并行模式运行。

由于这种方法结合了嵌入的模型并行性和多层感知器( MLP )的数据并行性,因此被称为混合并行训练(图 3 )。

pYYBAGJeYNWAdOnZAABU6rSXHbE578.png

图 3 。训练大型推荐系统的通用混合并行方法

Horovod all-to-all

从模型并行到数据并行范式需要一个多 GPU 集体通信操作:全部对全部。

All to All 是一种灵活的集体通信原语,可在每对 GPU 之间交换数据。这是必需的,因为在嵌入查找阶段结束时,每个 GPU 都保存所有样本的查找结果。但是,仅适用于表的子集(用于按表拆分)或列的子集(用于按列拆分)。

由于 all-to-all 操作会在 GPU 之间洗牌数据,因此需要注意的是,每个 GPU 都会保存所有表的所有列的嵌入查找结果,但只保存样本子集的嵌入查找结果。例如,对于一个 8 GPU 场景,本地批量大小毕竟是之前的 8 倍。

通信由 Horovod 库的 hvd.alltoall 函数处理。在引擎盖下,霍洛伍德称 NCCL 实施 为了获得最佳性能。如果你的系统上有 NVLink ,它也会利用它。

TensorFlow 2 中的混合并行训练示例

在本节中,我将描述一种用于 TensorFlow 2 中训练的 1130 亿参数推荐系统的混合并行训练方法。完整的源代码可以在 NVIDIA 深度学习示例库 中找到。

深度学习推荐模型的体系结构

对于这个例子,我使用 DLRM 体系结构(图 4 )。 DLRM 是研究论文 面向个性化和推荐系统的深度学习推荐模型 中首次介绍的一类推荐模型。我之所以选择它,是因为 MLPerf 基准测试使用了更小版本的 DLRM ,因此,它是演示推荐系统性能的当前行业标准。

DLRM 同时使用分类和数字功能。分类特征被输入到嵌入层中,而数字特征则由一个小的 MLP 子网络处理。

然后将这些层的结果输入点交互层和另一个 MLP 。然后使用二元交叉熵损失函数通过反向传播对模型进行训练,并根据随机梯度下降( SGD )方法更新权重。

poYBAGJeYNaAQs8jAAGCRawZKnw584.png

图 4 。 DLRM 体系结构图。

修改以支持宽深度模型

虽然我选择在本例中使用 DLRM 体系结构,但也可以支持相关模型(如 Wide & Deep )。这需要进行以下修改:

添加 wide & Deep 的“ wide ”部分,并在纯数据并行模式下运行它,完全绕过 all to all 。

为宽部分添加第二个优化器。

在深部,移除底部 MLP ,并将数字特征直接传递到顶部 MLP 。

移除点交互层。

同步文件夹

DLRM 可以在由数字和分类特征组成的任何表格数据集上进行训练。在本例中,我使用 Criteo 的 TB 点击日志数据集 ,因为它是最大的公开点击率数据集。

该数据集由 26 个分类变量和 13 个数值变量组成。在未经处理的数据中,独特类别的总数为 8.82 亿,其中 2.92 亿是在最大的特征中发现的。

遵循 MLPerf 推荐基准,对嵌入使用单精度,每个特征的嵌入维度为 128 。这意味着参数总数为 882M × 128 = 1130 亿。所有 26 个表的总大小为 1130 亿× 4 字节/ 230= 421 GiB ,最大表为 139.6 GiB 。因为最大的表不适合单个 GPU ,所以必须使用按列拆分模式将表分片,并将每个表分布到多个 GPU 中。

从理论上讲,您可以只对超过单个 GPU 内存的少数表执行此操作,并对其余的表使用按表拆分。然而,这将不必要地使代码复杂化,而没有任何明显的好处。因此,对所有表使用按列拆分模式。

性能优化

为了提高训练速度,我的团队实施了以下性能优化,如代码所示。这些是可以应用于其他深度学习推荐系统以及其他深度学习框架的通用策略。

自动混合精度

混合精度是计算方法中不同数值精度的组合使用。有关如何启用它的更多信息,请参阅 TensorFlow 核心文档中的 Mixed precision 。与 A100 上默认的 TF32 精度相比,该模型使用混合精度使其速度提高了 23% 。

相同宽度的融合嵌入表

当多个嵌入表具有相同的向量大小时——这是 DLRM 中使用embedding_dim=128的情况——它们可以沿零轴连接。这允许对一个大表执行单个查找,而不是对许多较小的表执行多个查找。

启动一个大内核而不是多个小内核要高效得多。在本例中,将表连接起来可使训练速度提高 39% 。

XLA

我的团队使用 TensorFlow 加速线性代数( XLA )编译器来提高性能。对于这个特定的用例,应用 XLA 比不使用它产生 3.36X 的加速。这个值是在打开所有其他优化的情况下实现的: AMP 、串联嵌入等等。

广播数据加载器

在每个 GPU 上运行每个嵌入表的一部分意味着每个 GPU 必须访问每个训练样本的每个特性。在每个过程中分别加载和解析所有这些输入数据效率低下,可能会导致严重的瓶颈。我通过只在第一个 worker 上加载输入数据并通过 NVLink 将其广播给其他 worker 来解决这个问题。这提供了 32% 的加速。

把这一切放在一起

图 5 显示了具有八个 GPU 的混合并行 DLRM 的设备放置示例。该图显示 GPU 0 和 7 。为了简单起见,它只显示分类功能 0 和 25 。

pYYBAGJeYNiAfd53AAWFuCOZ9AY516.png

图 5 。具有 1130 亿个参数的混合并行 DLRM 的简化图。

替代方法:将大型嵌入存储在 CPU 上

存储大型嵌入矩阵的一个简单替代方法是将它们放入主机内存中。小型嵌入表和计算密集型层仍然可以放置在 GPU 上,以获得最佳性能。虽然简单得多,但与将所有变量保留在 GPU 上相比,这种方法也较慢。

这有两个根本原因:

嵌入查找是一种内存受限的操作。 CPU 内存比 GPU 内存慢得多。双插槽 AMD Epyc 7742 的总内存带宽为 409.6 GB / s ,而单插槽 A100-80GB GPU 的总内存带宽为 2 TB / s ,而 8 个 A100-80GB GPU 的总内存带宽为 16 TB / s 。

GPU 之间的数据交换速度明显快于 CPU 和 GPU 之间的数据交换速度。这是因为将 CPU 连接到 GPU 之间的 PCIe 链路可能会成为瓶颈。

当使用 CPU 存储嵌入时, CPU 和 GPU 之间的传输必须首先通过提供 31.5 GB / s 带宽的 PCIe 接口。相反,在混合并行范例中,嵌入查找的结果通过 GPU 之间的 NVSwitch 结构进行传输。 DGX A100 采用第二代 NVSwitch 技术,支持每秒 600 GB 的峰值 GPU 到 – GPU 通信。

尽管速度有所放缓,但这种替代方法仍然比仅在 CPU 上运行整个网络快得多。

基准结果

下表显示了训练 113B 参数 DLRM 模型的基准测试结果。它只比较了三种硬件设置: CPU ,一种使用 CPU 内存的单一 GPU 用于最大的嵌入表,以及一种使用完整 DGX A100-80GB 的混合并行方法。

表 2 。比较 1130 亿参数深度学习推荐模型( DLRM )的 CPU 和 GPU 训练吞吐量。

比较前两行,你可以看到用一个 A100 GPU 来补充两个 CPU 可以使吞吐量增加 43 倍。之所以会出现这种情况,是因为 GPU 非常适合运行计算密集型线性层和适合其 80-GB 内存的较小嵌入层。

此外,使用八个 GPU 的完整 DGX A100 比在单个 A100 GPU 上训练快 15.5 倍。 DGX A100 使您能够将整个型号安装到 GPU 内存中,并消除了昂贵的设备到主机和主机到设备传输的需要。

总的来说, DGX A100 解决这项任务的速度是双插座 CPU 系统的 672 倍。

结论

在这篇文章中,我介绍了使用混合并行来训练大型推荐系统的想法。测试结果表明, DGX A100 是在 TensorFlow 2 中训练参数超过 1000 亿的推荐系统的极好工具。它在双插槽 CPU 上实现了 672 倍的加速。

高内存带宽和快速的 GPU 到 – GPU 通信使快速培训推荐人成为可能。因此,与仅使用 CPU 服务器相比,您的培训时间更短。这降低了培训成本,同时为从业者提供了更快的实验。

关于作者

Tomasz Grel 是一名深度学习工程师。在NVIDIA ,他专注于确保众多推荐系统的质量和执行速度,包括 NCF 、 VAE-CF 和 DLRM 。

审核编辑:郭婷

声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉
  • NVIDIA
    +关注

    关注

    14

    文章

    4958

    浏览量

    102882
  • gpu
    gpu
    +关注

    关注

    28

    文章

    4711

    浏览量

    128801
  • 深度学习
    +关注

    关注

    73

    文章

    5495

    浏览量

    121052
收藏 人收藏

    评论

    相关推荐

    芯青年,新征程——中科亿海微2024年度新员工培训顺利开班

    员工加入中科亿海微,随后从行业现状与产业突围、中科亿海微的创业实践、中科亿海微的文化构建以及如何做好工作等四个方面进行系统讲解,对我司中高
    的头像 发表于 07-18 08:12 553次阅读
    芯青年,新征程——中科<b class='flag-5'>亿</b>海微2024年度新员工<b class='flag-5'>培训</b>顺利开班

    何在Tensorflow实现反卷积

    ,扮演着重要角色。以下将详细阐述如何在TensorFlow实现反卷积,包括其理论基础、TensorFlow的实现方式、以及实际应用
    的头像 发表于 07-14 10:46 566次阅读

    TensorFlow是什么?TensorFlow怎么用?

    TensorFlow是由Google开发的一个开源深度学习框架,它允许开发者方便地构建、训练和部署各种复杂的机器学习模型。TensorFlow凭借其高效的计算性能、灵活的架构以及丰富的工具和库,在学
    的头像 发表于 07-12 16:38 611次阅读

    tensorflow和pytorch哪个更简单?

    工业界广泛使用、具有丰富生态系统和跨平台支持的框架,TensorFlow可能更适合您。以下是tensorflow和pytorch的介绍: TensorFlow和PyTorch的基本概念
    的头像 发表于 07-05 09:45 817次阅读

    tensorflow和pytorch哪个好

    。 生态系统 :拥有庞大的社区和丰富的资源,包括TensorFlow Hub、TensorBoard等工具。 易用性 :对于初学者来说,TensorFlow的学习曲线可能相对较陡
    的头像 发表于 07-05 09:42 651次阅读

    tensorflow简单的模型训练

    在本文中,我们将详细介绍如何使用TensorFlow进行简单的模型训练。TensorFlow是一个开源的机器学习库,广泛用于各种机器学习任务,包括图像识别、自然语言处理等。我们将从安装
    的头像 发表于 07-05 09:38 548次阅读

    keras模型转tensorflow session

    在这篇文章,我们将讨论如何将Keras模型转换为TensorFlow session。 Keras和TensorFlow简介 Keras是一个高级神经网络API,它提供了一种简单、快速的方式来构建
    的头像 发表于 07-05 09:36 483次阅读

    何在TensorFlow构建并训练CNN模型

    TensorFlow构建并训练一个卷积神经网络(CNN)模型是一个涉及多个步骤的过程,包括数据预处理、模型设计、编译、训练以及评估。下面,我将详细阐述这些步骤,并附上一个完整的代码示例。
    的头像 发表于 07-04 11:47 785次阅读

    TensorFlow的定义和使用方法

    数据流图,从而简化机器学习模型的构建、训练和部署。自2015年11月开源以来,TensorFlow迅速成为数据科学家、软件开发者以及教育工作者广泛使用的工具,广泛应用于图像识别、自然语言处理、推荐系统等多个领域。本文将深入解读Tenso
    的头像 发表于 07-02 14:14 720次阅读

    培训进行!米尔与瑞萨基于RZ/G2L的OpenAMP混合部署实战培训

    5月22-23日,由米尔电子和瑞萨电子联合举办的“瑞萨电子基于RZ/G2L的OpenAMP混合部署实战培训”在深圳福田区深业上城举办,课堂氛围非常活跃,下面我们一起来看看这两天培训现场的精彩瞬间吧
    的头像 发表于 05-30 08:01 408次阅读
    <b class='flag-5'>培训</b>进行<b class='flag-5'>中</b>!米尔与瑞萨基于RZ/G<b class='flag-5'>2</b>L的OpenAMP混合部署实战<b class='flag-5'>培训</b>

    双节锂电池拓扑交叉充电,主动平衡充电——ZCC1130T

    ZCC1130T是一款采用先进恒定电流/恒定电压算法的双节锂电池充电器,以其卓越性能和智能设计成为您设备的最佳搭档。最大支持1A的充电电流,确保您的电池迅速充满,时刻保持高效运行。 智能交叉充电
    发表于 01-08 15:55

    如何使用TensorFlow构建机器学习模型

    在这篇文章,我将逐步讲解如何使用 TensorFlow 创建一个简单的机器学习模型。
    的头像 发表于 01-08 09:25 935次阅读
    如何使用<b class='flag-5'>TensorFlow</b>构建机器学习模型

    微软宣布推出一个27亿参数的语言模型Phi-2

    Phi-2 是微软 “小语言模型(SLM)” 系列的最新版本。第一个版本是拥有 13 亿参数的 Phi-1,针对基本的 Python 编码任务进行了微调。
    发表于 12-15 09:46 258次阅读
    微软宣布推出一个27<b class='flag-5'>亿</b><b class='flag-5'>参数</b>的语言模型Phi-<b class='flag-5'>2</b>

    【飞腾派4G版免费试用】第二章:在PC端使用 TensorFlow2 训练目标检测模型

    使用 TensorFlow2 训练目标检测模型 因为我的项目是计划在飞腾派上实现一个目标检测跟踪算法,通过算法输出控制信号控制电机跟随目标运行。在第一章完成了Ubuntu系统的构建和烧写,这几天就在
    发表于 12-15 06:44

    【飞腾派4G版免费试用】 第二章:在PC端使用 TensorFlow2 训练目标检测模型

    本章记录了如何使用TensorFlow2 进行目标检测模型训练的过程。
    的头像 发表于 12-15 06:40 2517次阅读
    【飞腾派4G版免费试用】 第二章:在PC端使用 <b class='flag-5'>TensorFlow2</b> 训练目标检测模型