关于美团:TensorFlow在美团外卖推荐场景的GPU训练优化实践

美团机器学习平台基于外部深度定制的 TensorFlow 研发了 Booster GPU 训练架构。该架构在整体设计上充分考虑了算法、架构、新硬件的个性，从数据、计算、通信等多个角度进行了深度的优化，最终其性价比达到 CPU 工作的 2~4 倍。本文次要讲述 Booster 架构的设计实现、性能优化及业务落地工作，心愿能对从事相干开发的同学有所帮忙或者启发。

1 背景

在举荐零碎训练场景中，美团外部深度定制的 TenorFlow（简称 TF）版本[1]，通过 CPU 算力撑持了美团外部大量的业务。但随着业务的倒退，模型单次训练的样本量越来越多，构造也变得越来越简单。以美团外卖举荐的精排模型为例，单次训练的样本量已达百亿甚至千亿，一次试验要消耗上千核，且优化后的训练任务 CPU 使用率已达 90% 以上。为了反对业务的高速倒退，模型迭代试验的频次和并发度都在一直减少，进一步减少了算力应用需要。在估算无限的前提下，如何以较高的性价比来实现高速的模型训练，从而保障高效率的模型研发迭代，是咱们迫切需要解决的问题。

近几年，GPU 服务器的硬件能力突飞猛进，新一代的 NVIDIA A100 80GB SXM GPU 服务器（8 卡）[2]，在存储方面能够做到：显存 640GB、内存 1~2TB、SSD10+TB，在通信方面能够做到：卡间双向通信 600GB/s、多机通信 800~1000Gbps/s，在算力方面能够做到：GPU 1248TFLOPS（TF32 Tensor Cores），CPU 96~128 物理核。如果训练架构能充分发挥新硬件的劣势，模型训练的老本将会大大降低。但 TensorFlow 社区在举荐零碎训练场景中，并没有高效和成熟的解决方案。咱们也尝试应用优化后的 TensorFlow CPU Parameter Server[3]（简称 PS）+GPU Worker 的模式进行训练，但其只对简单模型有肯定的收益。NVIDIA 开源的 HugeCTR[4]尽管在经典的深度学习模型上性能体现优异，但要在美团的生产环境间接应用起来，还须要做较多的工作。

美团根底研发机器学习平台训练引擎团队，联结到家搜推技术部算法效力团队、NVIDIA DevTech 团队，成立了联结项目组。在美团外部深度定制的 TenorFlow 以及 NVIDIA HugeCTR 的根底上，研发了 举荐零碎场景的高性能 GPU 训练架构 Booster。目前在美团外卖举荐场景中进行了部署，多代模型全面对齐算法的离线成果，比照之前，优化后的 CPU 工作，性价比晋升了 2~4 倍。因为 Booster 对原生 TensorFlow 接口有较好的兼容性，原 TensorFlow CPU 工作只须要一行代码就可实现迁徙。这样让 Booster 能够疾速在美团多条业务线上进行初步验证，相比之前的 CPU 工作，均匀性价比都晋升到 2 倍以上。本文将重点介绍 Booster 架构的设计与优化，以及在美团外卖举荐场景落地的全过程，心愿能对大家有所帮忙或启发。

2 GPU 训练优化挑战

GPU 训练在美团内曾经广泛应用到 CV、NLP、ASR 等场景的深度学习模型，但在举荐零碎场景中，却迟迟没有失去大规模的利用，这跟场景的模型特点、GPU 服务器的硬件特点都有较强的关系。

举荐零碎深度学习模型特点

• 读取样本量大：训练样本在几十 TB~ 几百 TB，而 CV 等场景通常在几百 GB 以内。
• 模型参数量大：同时有大规模稠密参数和浓密参数，须要几百 GB 甚至上 TB 存储，而 CV 等场景模型次要是浓密参数，通常在几十 GB 以内。
• 模型计算复杂度绝对低一些：举荐零碎模型在 GPU 上单步执行只须要 10~100ms，而 CV 模型在 GPU 上单步执行是 100~500ms，NLP 模型在 GPU 上单步执行是 500ms~1s。

GPU 服务器特点

• GPU 卡算力很强，但显存仍无限 ：如果要充分发挥 GPU 算力，须要把 GPU 计算用到的各种数据提前搁置到显存中。而从 16 年~20 年，NVIDIA Tesla GPU 卡[5] 计算能力晋升了 10 倍以上，但显存大小只晋升了 3 倍左右。
• 其它维度资源并不是很短缺：相比 GPU 算力的晋升速度，单机的 CPU、网络带宽的增长速度较慢，如果遇到这两类资源 workload 较重的模型，将无奈充分发挥 GPU 的能力，GPU 服务器相比 CPU 服务器的性价比不会太高。

总结来说，CV、NLP 等场景的模型训练属于计算密集型工作，而且大多模型单张卡的显存都能够装下，这和 GPU 服务器的劣势十分好地进行了匹配。但在举荐零碎场景中，因为模型绝对没有那么简单，远端读取的样本量大，特色解决消耗 CPU 多，给单机 CPU 和网络带来较大的压力。同时面对模型参数量大的状况，单机的 GPU 显存是无奈放下的。这些 GPU 服务器的劣势，恰好都被举荐零碎场景命中。

好在 NVIDIA A100 GPU 服务器，在硬件上的降级补救了显存、CPU、带宽这些短板，但如果零碎实现和优化不当，仍然不会有太高的性价比收益。在落地 Booster 架构的过程中，咱们次要面临如下挑战：

• 数据流零碎：如何利用好多网卡、多路 CPU，实现高性能的数据流水线，让数据的供应能够跟上 GPU 的生产速度。
• 混合参数计算：对于大规模稠密参数，GPU 显存间接装不下的状况，如何充分利用 GPU 高算力、GPU 卡间的高带宽，实现一套大规模稠密参数的计算，同时还须要兼顾浓密参数的计算。

3 零碎设计与实现

面对下面的挑战，如果纯从零碎的的角度去设计，难度较大。Booster 采纳了“算法 + 零碎”Co-design 的设计思路，让这代零碎的设计大大失去简化。在系统实施门路上，思考到业务预期交付工夫、施行危险，咱们并没有一步到位落地 Booster 的多机多卡版本，而是第一版先落地了 GPU 单机多卡版本，本文重点介绍的也是单机多卡的工作。另外，依靠于 NVIDIA A100 GPU 服务器弱小的计算能力，单机的算力能够满足美团绝大多数业务的单次试验需要。

3.1 参数规模的合理化

大规模稠密离散特色的应用，导致深度预估模型的 Embedding 参数量急剧收缩，数 TB 大小的模型一度风行于业界推搜的各大头部业务场景。然而业界很快意识到，在硬件老本无限的状况下，过于宏大的模型给生产部署运维和试验迭代翻新削减了惨重的累赘。学术研究表明[10-13]，模型成果强依赖于模型的信息容量，并非参数量。实践证明，前者能够通过模型构造的优化来进行晋升，而后者在保障成果的前提下，尚存有很大的优化空间。Facebook 在 2020 年提出了 Compositional Embedding[14]，实现举荐模型参数规模数个量级的压缩。阿里巴巴也发表了相干工作[15]，将外围业务场景的预估模型由数 TB 压缩至几十 GB 甚至更小。总的来看，业界的做法次要有以下几种思路：

• 去穿插特色 ：穿插特色由单特色间做笛卡尔积产生，这会生成微小的特色 ID 取值空间和对应 Embedding 参数表。深度预估模型倒退至今，曾经有大量的办法通过模型构造来建模单特色间的交互，防止了穿插特色造成的 Embedding 规模收缩，如 FM 系列[16]、AutoInt[17]、CAN[18] 等。
• 精简特色 ：特地是基于 NAS 的思路，以较低的训练老本实现深度神经网络自适应特征选择，如 Dropout Rank[19] 和 FSCD[20]等工作。
• 压缩 Embedding 向量数 ：对特色取值进行复合 ID 编码和 Embedding 映射，以远小于特色取值空间的 Embedding 向量数，来实现丰盛的特色 Embedding 表白，如 Compositional Embedding[14]、Binary Code Hash Embedding[21] 等工作。
• 压缩 Embedding 向量维度 ：一个特色 Embedding 向量的维度决定了其表征信息的下限，然而并非所有的特色取值都有那么大的信息量，须要 Embedding 表白。因而，能够每一个特征值自适应的学习精简 Embedding 维度，从而压缩参数总量，如 AutoDim[22] 和 AMTL[23]等工作。
• 量化压缩 ：应用半精度甚至 int8 等更激进的形式，对模型参数做量化压缩，如 DPQ[24] 和 MGQE[25]。

美团外卖举荐的模型一度达到 100G 以上，通过利用以上计划，咱们在模型预估精度损失可控的前提下，将模型管制在 10GB 以下。

基于这个算法根底假如，咱们将 第一阶段的设计指标定义到反对 100G 以下的参数规模。这能够比拟好的适配 A100 的显存，寄存在单机多卡上，GPU 卡间双向带宽 600GB/s，能够充分发挥 GPU 的解决能力，同时也能够满足美团大多数模型的需要。

3.2 零碎架构

基于 GPU 零碎的架构设计，要充分考虑硬件的个性能力充分发挥性能的劣势。咱们 NVIDIA A100 服务器的硬件拓扑和 NVIDIA DGX A100[6]比拟相似，每台服务器蕴含：2 颗 CPU，8 张 GPU，8 张网卡。Booster 架构的架构图如下所示：

整个零碎次要包含三个外围模块：数据模块，计算模块，通信模块：

• 数据模块：美团自研了一套反对多数据源、多框架的数据散发零碎，在 GPU 零碎上，咱们革新数据模块反对了多网卡数据下载，以及思考到 NUMA Awareness 的个性，在每颗 CPU 上都部署了一个数据散发服务。
• 计算模块：每张 GPU 卡启动一个 TensorFlow 训练过程执行训练。
• 通信模块 ：咱们应用了 Horovod[7] 来做分布式训练的卡间通信，咱们在每个节点上启动一个 Horovod 过程来执行对应的通信工作。

上述的设计，合乎 TensorFlow 和 Horovod 原生的设计范式。几个外围模块能够互相解耦，独立迭代，而且如果合并开源社区的最新个性，也不会对系统造成架构性的冲击。

咱们再来看一下整个零碎的简要执行流程，每张 GPU 卡上启动的 TensorFlow 过程外部的执行逻辑如下图：

整个训练流程波及参数存储、优化器、卡间通信等几个要害模块。对于样本的输出特色，咱们分为稠密特色（ID 类特色）和浓密特色。在理论业务场景中，稠密特色通常 IDs 总量较多，对应的稠密参数应用 HashTable 数据结构存储更适合，而且因为参数量较大，GPU 单卡显寄存不下，咱们会通过 ID Modulo 的形式 Partition 到多张 GPU 卡的显存中寄存。对于 IDs 总量较少的稠密特色，业务通常应用多维矩阵数据结构表白（在 TensorFlow 外面的数据结构是 Variable），因为参数量不大，GPU 单卡显存能够放下，咱们应用 Replica 的形式，每张 GPU 卡的显存都搁置一份参数。对于浓密参数，通常应用 Variable 数据结构，以 Replica 的形式搁置到 GPU 显存中。下边将具体介绍 Booster 架构的外部实现。

3.3 要害实现

3.3.1 参数存储

早在 CPU 场景的 PS 架构下，咱们就实现了大规模稠密参数的整套逻辑，当初要把这套逻辑搬到 GPU 上，首先要实现的就是 GPU 版本的 HashTable。咱们调研了业界多种 GPU HashTable 的实现，如 cuDF、cuDPP、cuCollections、WarpCore 等，最终抉择了基于 cuCollections 实现 TensorFlow 版本的 GPUHashTable。究其原因，次要是因为理论业务场景中，大规模稠密特色的总量通常是未知的，并且随时可能呈现特色穿插，从而以致稠密特色的总量变化很大，这就导致“动静扩容”能力将成为咱们 GPU HashTable 的必备性能，可能做到动静扩容的只有 cuCollections 的实现。咱们在 cuCollections 的 GPU HashTable 根底上实现了非凡接口（find\_or\_insert），对大规模读写性能进行了优化，而后封装到了 TensorFlow 中，并在其上实现了低频过滤的性能，能力上对齐 CPU 版本的稠密参数存储模块。

3.3.2 优化器

目前，稠密参数的优化器与浓密参数的优化器并不兼容，咱们在 GPU HashTable 的根底上，实现了多种稠密优化器，并且都做了优化器动量 Fusion 等性能，次要实现了 Adam、Adagrad、FTRL、Momentum 等优化器。对理论业务场景来说，这些优化器曾经可能笼罩到绝大多数业务的应用。浓密局部参数能够间接应用 TensorFlow 原生反对的稠密 / 浓密优化器。

3.3.2 卡间通信

理论训练期间，对于不同类型的特色，咱们的解决流程也有所不同：

• 稠密特色（ID 类特色，规模较大，应用 HashTable 存储）：因为每张卡的输出样本数据不同，因而输出的稠密特色对应的特征向量，可能寄存在其余 GPU 卡上。具体流程上，训练的前向咱们通过卡间 AllToAll 通信，将每张卡的 ID 特色以 Modulo 的形式 Partition 到其余卡中，每张卡再去卡内的 GPUHashTable 查问稠密特征向量，而后再通过卡间 AllToAll 通信，将第一次 AllToAll 从其余卡上拿到的 ID 特色以及对应的特征向量原路返回，通过两次卡间 AllToAll 通信，每张卡样本输出的 ID 特色都拿到对应的特征向量。训练的反向则会再次通过卡间 AllToAll 通信，将稠密参数的梯度以 Modulo 的形式 Partition 到其余卡中，每张卡拿到本人的稠密梯度后再执行稠密优化器，实现大规模稠密特色的优化。具体流程如下图所示：

• 稠密特色（规模较小，应用 Variable 存储）：相比应用 HashTable 的区别，因为每张 GPU 卡都有全量的参数，间接在卡内查找模型参数即可。在反向聚合梯度的时候，会通过卡间 AllGather 获取所有卡上的梯度求均匀，而后交给优化器执行参数优化。
• 浓密特色：浓密参数也是每张卡都有全量的参数，卡内能够间接获取参数执行训练，最初通过卡间 AllReduce 聚合多卡的浓密梯度，执行浓密优化器。

在整个的执行过程中，稠密参数和浓密参数全副搁置在 GPU 显存中，模型计算也全副在 GPU 上解决，GPU 卡间通信带宽也足够快，可能充分发挥了 GPU 的弱小算力。

这里小结一下，Booster 训练架构，与 CPU 场景 PS 架构的外围区别在于：

• 训练模式：PS 架构是异步训练模式，Booster 架构是同步训练模式。
• 参数散布：PS 架构下模型参数都寄存在 PS 内存中，Booster 架构下稠密参数（HashTable）是 Partition 形式散布在单机八卡中，浓密参数（Variable）是 Replica 形式寄存在每张卡中，因而 Booster 架构下的 Worker 角色兼顾了 PS 架构下 PS/Worker 角色的性能。
• 通信形式：PS 架构下 PS/Worker 间通信走的是 TCP（Grpc/Seastar），Booster 架构下 Worker 间通信走的是 NVSwitch（NCCL），任意两卡间双向带宽 600GB/s，这也是 Booster 架构的训练速度获得较大晋升的起因之一。

因为每张卡的输出数据不同，并且模型参数既有在卡间 Partition 存储的，也有在卡间 Replica 存储的，因而 Booster 架构同时存在模型并行、数据并行。此外，因为 NVIDIA A100 要求 CUDA 版本 >=11.0，而 TensorFlow 1.x 版本只有 NV1.15.4 才反对 CUDA11.0。美团绝大多数业务场景都还在应用 TensorFlow 1.x，因而咱们所有革新都是在 NV1.15.4 版本根底上开发的。

以上就是 Booster 整体零碎架构及外部执行流程的介绍。下文次要介绍在初步实现的 Booster 架构的根底上，咱们所做的一些性能优化工作。

4 零碎性能优化

基于上述的设计实现完第一版零碎后，咱们发现端到端性能并不是很合乎预期，GPU 的 SM 利用率（SM Activity 指标）只有 10%~20%，相比 CPU 并没有太大的劣势。为了剖析架构的性能瓶颈，咱们应用 NVIDIA Nsight Systems（以下简称 nsys）、Perf、uPerf 等工具，通过模块化压测、模仿剖析等多种剖析伎俩，最终定位到数据层、计算层、通信层等几方面的性能瓶颈，并别离做了相应的性能优化。以下咱们将以美团外卖某举荐模型为例，别离从 GPU 架构的数据层、计算层、通信层，一一介绍咱们所做的性能优化工作。

4.1 数据层

如前文所述，举荐零碎的深度学习模型，样本量大，模型绝对不简单，数据 I / O 自身就是瓶颈点。如果几十台 CPU 服务器上的数据 I / O 操作，都要在单台 GPU 服务器上实现，那么数据 I / O 的压力会变得更大。咱们先看一下在以后零碎下的样本数据流程，如下图所示：

外围流程：数据散发过程通过网络读取 HDFS 样本数据（TFRecord 格局）到内存中，而后通过共享内存（Shared Memory）的形式把样本数据传输给 TensorFlow 训练过程。TensrFlow 训练过程收到样本数据后，走原生的 TensrFlow 特色解析逻辑，拿到特色数据后通过 GPU MemcpyH2D 到 GPU 显存中。咱们通过模块化压测剖析发现，数据散发层的样本拉取、TensrFlow 层的特色解析以及特色数据 MemcpyH2D 到 GPU 等几个流程，都存在较大的性能问题（图中黄色流程所示），以下具体介绍咱们在这几块所做的性能优化工作。

4.1.1 样本拉取优化

样本拉取、组装 Batch 是由数据散发过程实现的，咱们在这里所做的次要优化工作是，首先将数据散发过程通过 numactl 独立到 NUMA 外部执行，防止了 NUMA 间的数据传输；其次，数据下载从单网卡裁减到了多网卡，增大数据下载带宽；最初，数据散发过程与 TensrFlow 过程之间的传输通道，从单个 Shared Memory 扩大到每张 GPU 卡有独立的 Shared Memory，防止了单 Shared Memory 所带来的内存带宽问题，并在 TensrFlow 外部实现了特色解析时对输出数据零拷贝的能力。

4.1.2 特色解析优化

目前，美团外部绝大多数业务的样本数据都还是 TFRecord 格局，TFRecord 实际上是 ProtoBuf（简称 PB）格局。PB 反序列化十分消耗 CPU，其中 ReadVarint64Fallback 办法 CPU 占用较为突出，理论 profiling 后果如下图：

究其原因，CTR 场景的训练样本通常蕴含了大量的 int64 类型的特色，int64 在 PB 中是以 Varint64 类型数据存储的，ReadVarint64Fallback 办法就是用来解析 int64 类型的特色。一般的 int64 数据类型须要占用 8 个字节，而 Varint64 针对不同的数据范畴，应用了变长的存储长度。PB 在解析 Varint 类型数据时，首先要确定以后数据的长度，Varint 用 7bit 存储数据，高位 1bit 存储标记位，该标记位示意下一个字节是否无效，如果以后字节最高位为 0，则阐明以后 Varint 数据在该字节处完结。咱们理论业务场景的 ID 特色大多是通过 Hash 后的值，用 Varint64 类型表白会比拟长，这也就导致在特色解析过程中要屡次判断数据是否完结，以及屡次位移和拼接来生成最终数据，这使得 CPU 在解析过程中存在大量的 分支预测和长期变量，十分影响性能。以下是 4 字节 Varint 的解析流程图：

这个解决流程，非常适合用 SIMD 指令集批处理优化。以 4 字节的 Varint 类型为例，咱们的优化流程次要包含两步：

1. SIMD 寻找最高位：通过 SIMD 指令将 Varint 类型数据的每个字节与 0xF0 做与运算，找到第一个后果等于 0 的字节，这个字节就是以后 Varint 数据的完结地位。
2. SIMD 解决 Varint：按理来说，通过 SIMD 指令将 Varint 数据高位清零后的每个字节顺次右移 3 /2/1/ 0 字节，就可失去最终的 int 类型数据，但 SIMD 没有这样的指令。因而，咱们通过 SIMD 指令别离解决每个字节的高 4bit、低 4bit，实现了这个性能。咱们将 Varint 数据的高下 4bit 别离解决成 int\_h4 与 int\_l4，再做或运算，就失去了最终的 int 类型数据。具体优化流程如下图所示（4 字节数据）：

对于 Varint64 类型数据的解决，咱们间接分成了两个 Varint 类型数据来解决。通过这两步的 SIMD 指令集优化，样本解析速度失去大大晋升，在 GPU 端到端训练速度晋升的同时，CPU 使用率降落了 15%。这里咱们次要应用了 SSE 指令集优化，期间也尝试了 AVX 等更大长度的指令集，但成果不是很显著，最终并没有应用。此外，SIMD 指令集在老的机器上会导致 CPU 重大降频，因而官网社区并没有引入这个优化，而咱们 GPU 机器的 CPU 都比拟新，齐全能够应用 SIMD 指令集进行优化。

4.1.3 MemcpyH2D 流水线

解析完样本失去特色数据后，须要将特色数据拉到 GPU 中能力执行模型计算，这里须要通过 CUDA 的 MemcpyH2D 操作。咱们通过 nsys 剖析这块的性能，发现 GPU 在执行期间有较多的进展工夫，GPU 须要期待特色数据 Memcpy 到 GPU 上之后能力执行模型训练，如下图所示：

对于 GPU 零碎的数据流，须要提前传输到离 GPU 处理器最近的显存中，能力施展 GPU 的计算能力。咱们基于 TensorFlow 的 prefetch 性能，实现了 GPU 版本的 PipelineDataset，在计算之前先把数据拷贝到了 GPU 显存中。须要留神的是 CPU 内存拷贝到 GPU 显存这个过程，CPU 内存须要应用 Pinned Memory，而非原生的 Paged Memory，能够减速 MemcpyH2D 流程。

4.1.4 硬件调优

在数据层的性能优化期间，美团外部根底研发平台的服务器组、网络组、操作系统组也帮忙咱们做了相干的调优：

• 在网络传输方面，为了缩小网络协议栈解决开销，进步数据拷贝的效率，咱们通过优化网卡配置，开启 LRO（Large-Receive-Offload）、TC Flower 的硬件卸载、Tx-Nocache-Copy 等个性，最终网络带宽晋升了 17%。
• 在 CPU 性能优化方面，通过性能 profiling 剖析，发现内存提早和带宽是瓶颈。于是咱们尝试了 3 种 NPS 配置，综合业务场景和 NUMA 个性，抉择了 NPS2。此外，联合其余 BIOS 配置（例如 APBDIS，P-state 等），能够将内存提早升高 8%，内存带宽晋升 6%。

通过上述优化，网络极限带宽晋升了 80%，在业务需要带宽下 GPU 的 H2D 带宽晋升了 86%。最终在数据解析层面也拿到了 10%+ 的性能收益。

通过数据层样本拉取、特色解析、MemcpyH2D 和硬件的优化，Booster 架构端到端训练速度晋升了 40%，训练性价比达到了 CPU 的 1.4 倍，数据层也不再成为以后架构的性能瓶颈。

4.2 计算层

4.2.1 Embedding 流水线

早在 CPU 场景做 TensorFlow 训练性能优化时，咱们就曾经实现了 Embedding Pipeline[1]的性能：咱们把整个计算图拆分为 Embedding Graph（EG）和 Main Graph（MG）两张子图，两者异步独立执行，做到执行上的 Overlap（整个拆分过程，能够做到对用户通明）。EG 次要笼罩从样本中抽取 Embedding Key，查问组装 Embedding 向量，Embedding 向量更新等环节；MG 次要蕴含浓密局部子网络计算、梯度计算、浓密参数局部更新等环节。

两张子图的交互关系为：EG 向 MG 传递 Embedding 向量（从 MG 的视角看，是从一个浓密 Variable 读取数值），MG 向 EG 传递 Embedding 参数对应的梯度。上述两个过程的表白都是 TensorFlow 的计算图，咱们利用两个 Python 线程，两个 TensorFlow Session 并发的执行两张计算图，使得两个阶段 Overlap 起来，以此达到了更大的训练吞吐。

咱们把这个流程在 GPU 架构下也实现了一遍，并在其中退出了卡间同步流程，大规模稠密特色的 AllToAll 通信及其反向梯度的 AllToAll 通信都在 EG 中执行，一般稠密特色的反向梯度的卡间 AllGather 同步、浓密参数的反向梯度的卡间 AllReduce 同步都在 MG 中执行。须要留神的是，在 GPU 场景中，EG、MG 是在同一个 GPU Stream 上执行 CUDA Kernel 的，咱们尝试过 EG、MG 别离在独立的 GPU Stream 上执行，性能会变差，深层起因与 CUDA 底层实现无关，这个问题自身还在期待解决。

4.2.2 算子优化及 XLA

相比 CPU 层面的优化，GPU 上的优化更加简单。首先对于 TensorFlow 的算子，还有一些没有 GPU 的实现，当模型中应用了这些 CPU 算子，会跟上下游的 GPU 算子呈现内存和显存之间的数据来回拷贝，影响整体性能，咱们在 GPU 上实现了应用较为频繁、影响较大的算子。另外，对于 TensorFlow 这代框架，算子粒度是十分细的，能够不便用户灵便搭建各种简单的模型，但这对 GPU 处理器来说却是一个劫难，大量的 Kernel Launch 以及访存开销导致不能充分利用 GPU 算力。对于 GPU 上的优化，通常有两个方向，手工优化和编译优化。在手工优化方面，咱们从新实现了一些罕用的算子和层（Unique、DynamicPartition、Gather 等）。

以 Unique 算子为例，原生 TensorFlow 的 Unique 算子要求输入元素的程序与输出元素的程序统一，而在理论场景中，咱们并不需要这个限度，咱们批改了 Unique 算子的 GPU 实现，缩小了因输入有序导致的额定执行的 GPU Kernel。

在编译优化方面，目前咱们次要应用 TensorFlow 社区提供的 XLA[9]来做一些主动优化。原生 TensorFlow 1.15 中的 XLA 失常开启可取得 10～20% 端到端的性能晋升。但 XLA 对算子动静 shape 不能很好地进行反对，而举荐零碎场景的模型中这种状况却十分常见，这就导致 XLA 减速性能不合乎预期，甚至是负优化，因而咱们做了如下的缓解工作：

• 局部优化：对于咱们手动引入的动静 shape 算子（如 Unique），咱们进行了子图标记，不执行 XLA 编译，XLA 只优化能够稳固减速的子图。
• OOM 兜底：XLA 会依据算子的 type、input type、shape 等信息，缓存编译两头后果，防止反复编译。然而因为稠密场景以及 GPU 架构实现的特殊性，人造存在 Unique、DynamicPartition 等 Output shape 是动静的算子，这就导致这些算子以及连贯在这些算子之后的算子，在执行 XLA 编译时无奈命中 XLA 缓存而从新编译，新的缓存越来越多，而旧的缓存不会被开释，最终导致 CPU 内存 OOM。咱们在 XLA 外部实现了 LRUCache，被动淘汰掉旧的 XLA 缓存，防止 OOM 的问题。
• Const Memcpy 打消：XLA 在应用 TF\_HLO 重写 TensorFlow 算子时，对一些编译期已固定的数据会打上 Const 标记，然而这些 Const 算子的 Output 只能定义在 Host 端，为了将 Host 端的 Output 送给 Device 端须要再加一次 MemcpyH2D，这就占用了 TensorFlow 原有的 H2D Stream，影响样本数据提前拷贝到 GPU 端。因为 XLA 的 Const Output 在编译期曾经固化，因而没有必要每一步都做一次 MemcpyH2D，咱们将 Device 端的 Output 缓存下来，后续应用该 Output 时，间接从缓存中读取，防止多余的 MemcpyH2D。

对于 XLA 的优化，确切的来说应该是问题修复，目前可能做到的是 GPU 场景下能够失常开启 XLA，并取得 10～20% 的训练速度晋升。值得一提的是，对于动静 shape 的算子编译问题，美团外部根底研发机器学习平台 / 深度学习编译器团队曾经有了彻底的解决方案，后续咱们会联结解决这个问题。

通过计算层的 Embedding 流水线、XLA 相干优化，Booster 架构端到端训练速度晋升了 60%，GPU 单机八卡训练性价比达到等同资源下 CPU 的 2.2 倍。

4.3 通信层

在单机多卡训练过程中，咱们通过 Nsight Systems 剖析发现，卡间通信耗时占比十分高，而且在此期间 GPU 使用率也非常低，如下图所示：

从图中能够看出，训练期间卡间通信耗时比拟长，同时在通信期间 GPU 使用率也非常低，卡间通信是影响训练性能晋升的要害瓶颈点。咱们对通信过程进行拆解打点后发现，卡间通信（AllToAll、AllReduce、AllGather 等）协商的工夫远远高于数据传输的工夫：

剖析具体起因，以负责大规模稠密参数通信的 AllToAll 为例，咱们通过 Nsight Systems 工具，察看到通信协商工夫长次要是因为某张卡上的算子执行工夫比拟晚导致的。因为 TensorFlow 算子调度并不是严格有序，同一个特色的 embedding\_lookup 算子，在不同卡上真正执行的工夫点也不尽相同，某张卡上第一个执行 embedding\_lookup 算子在另一张卡上可能是最初一个执行，因而咱们狐疑不同卡上算子调度的不一致性，导致了各张卡发动通信的时刻不同，并最终导致了通信协商工夫过长。咱们通过几组模拟实验也论证了的确是由算子调度导致的。对于这个问题，最间接的想法是革新 TensorFlow 计算图的外围调度算法，但这个问题在学术界也始终是一个简单的问题。咱们换了一种思路，通过交融要害的算子，来缓解这个问题，通过统计，咱们抉择了 HashTable 和 Variable 相干的算子。

4.3.1 HashTable 相干算子交融

咱们设计和实现了一个图优化过程，这个过程会主动地将图中能够合并的 HashTable 及对应的 embedding\_lookup 过程进行合并，合并策略上次要将 embedding\_size 雷同的 HashTable 合并到一块。同时为了防止 HashTable 合并之后原始特色之间产生 ID 抵触，咱们引入了主动对立特色编码的性能，对不同的原始特色别离加上不同的偏移量，纳入不同的特色域，实现了训练时的对立特色编码。

咱们在某理论业务模型上进行测试，该图优化将 38 张 HashTable 合并成为了 2 张 HashTable，将 38 次 embedding\_lookup 合并成了 2 次，这将 EmbeddingGraph 中的 embedding\_lookup 相干算子数量缩小了 90%，卡间同步通信次数缩小了 90%。此外，算子合并之后，embedding\_lookup 中的 GPU 算子也产生了合并，缩小了 Kernel Launch 次数，使得 EmbeddingGraph 的执行速度变得更快。

4.3.2 Variable 相干算子交融

相似于 HashTable Fusion 的优化思路，咱们察看到业务模型中通常蕴含数十至数百个 TensorFlow 原生的 Variable，这些 Variable 在训练期间梯度须要做卡间同步，同样的，Variable 数量太多导致卡间同步的协商工夫变长。咱们通过 Concat/Split 算子，将所有的 Trainable Variables 主动合并到一起，使得整个 MG 的反向只产生几个梯度 Tensor，大大减少了卡间同步的次数。同时，做完 Variable Fusion 之后，优化器中理论执行的算子数量也大大减少，放慢了计算图自身的执行速度。

须要留神的是，TensorFlow 的 Variable 分为两种，一种是每个 Step 全副参数值都参加训练的 Dense Variable，如 MLP 的 Weight；另一种是专门用于 embedding\_lookup 的 Variable，每个 Step 只有局部值参加训练，咱们称之为 Sparse Variable。对于前者，做 Variable 合并不会影响到算法成果。而对于后者，它反向梯度是 IndexedSlices 对象，卡间同步默认走的是 AllGather 通信，如果业务模型中对于 Sparse Variables 的优化采纳的是 Lazy 优化器，即每个 Step 只优化更新 Variable 中的某些行，此时对 Sparse Variables 做合并，会导致其反向梯度从 IndexedSlices 对象转为 Tensor 对象，卡间同步变成 AllReduce 过程，就可能会影响到算法成果。对于这种状况，咱们提供了一个开关，由业务去管制是否合并 Sparse Variables。通过咱们的实测，在某举荐模型上合并 Sparse Variables 会进步 5～10% 的训练性能，而对理论业务成果的影响在一个千分点以内。

这两种算子交融的优化，不仅优化了卡间通信性能，对卡内计算性能也有肯定的晋升。通过这两种算子交融的优化，GPU 架构端到端训练速度晋升了 85%，同时不影响业务算法的成果。

4.4 性能指标

实现了数据层、计算层、通信层的性能优化后，比照咱们的 TensorFlow[3] CPU 场景，GPU 架构获得了 2～4 倍的性价比收益（不同业务模型收益不同）。咱们基于美团外卖某举荐模型，应用单台 GPU 节点（A100 单机八卡）和同老本的 CPU Cluster，别离比照了原生 TensorFlow 1.15 和咱们优化后的 TensorFlow 1.15 的训练性能，具体数据如下：

能够看到，咱们优化后的 TensorFlow GPU 架构训练吞吐，是原生 TensorFlow GPU 的 3 倍以上，是优化后 TensorFlow CPU 场景的 4 倍以上。

注：原生 TensorFlow 应用了 tf.Variable 作为 Embedding 的参数存储。

5 业务落地

Booster 架构要在业务生产中落地，不只是要有一个良好的零碎性能，还须要同时关注训练生态系统的齐备性以及训练产出模型的成果。

5.1 齐备性

一次残缺的模型训练试验，除了要跑训练（Train）工作外，往往还须要跑模型的成果评估（Evaluate）或模型的预估（Predict）工作。咱们基于 TensorFlow Estimator 范式对训练架构进行封装，实现用户侧一套代码对立反对 GPU 和 CPU 场景下的 Train、Evaluate 和 Predict 工作，通过开关进行灵便切换，用户只须要关注模型代码自身的开发。咱们将架构改变全都封装到了引擎外部，用户只须要一行代码就能从 CPU 场景迁徙到 GPU 架构：

 tf.enable_gpu_booster()

理论业务场景，用户通常会应用 train\_and\_evaluate 模式，在跑训练任务的过程中同时评估模型成果。上了 Booster 架构后，因为训练跑的太快，导致 Evaluate 速度跟不上训练失常产出 Checkpoint 的速度。咱们在 GPU 训练架构的根底上，反对了 Evaluate on GPU 的能力，业务能够申请一颗 A100 GPU 专门用来做 Evaluate，单颗 GPU 做 Evaluate 的速度是 CPU 场景下单个 Evaluate 过程的 40 倍。同时，咱们也反对了 Predict on GPU 的能力，单机八卡 Predict 的速度是等同老本下 CPU 的 3 倍。

此外，咱们在工作资源配置上也提供了比较完善的选项。在单机八卡（A100 单台机器至少配置 8 张卡）的根底上，咱们反对了单机单卡、双卡、四卡工作，并买通了单机单卡 / 双卡 / 四卡 / 八卡 /CPU PS 架构的 Checkpoint，使得用户可能在这几种训练模式间自在切换、断点续训，不便用户抉择正当的资源类型、资源量跑试验，同时业务也可能从已有模型的 Checkpoint 来 WarmStart 训练新的模型。

5.2 训练成果

相较 PS/Worker 异步模式的 CPU 训练，单机多卡训练时卡间是全同步的，因此防止了异步训练梯度更新提早对训练成果的影响。然而，因为同步模式下每一步迭代的理论 Batch Size 是每张卡样本数的总和，并且为了充分利用 A100 卡的算力，咱们会将每张卡的 Batch Size（单步迭代的样本数）尽量调大。这使得理论训练的 Batch Size（1 万~10 万）比 PS/Worker 异步模式（1 千~1 万）大很多。咱们须要面临大 Batch 下训练超参调优的问题[26,27]：在保障 Epoch 不变的前提下，扩充 Batch Size 会导致参数无效更新次数缩小，可能导致模型训练的成果变差。

咱们采纳 Linear Scaling Rule[28]的准则领导调整学习率。如果训练 Batch Size 较 PS/Worker 模式的 Batch Size 增大 N 倍，将学习率也放大 N 倍即可。这种形式简略便于操作，实际成果还不错。当然须要留神的是，如果原有训练形式的学习率曾经很激进时，大 Batch Size 训练学习率的调整幅度则须要适当减小，或者应用学习率 Warmup 等更简单的训练策略[29]。咱们会在后续工作中对超参优化模式做更深刻的摸索。

6 总结与瞻望

在美团举荐零碎训练场景，随着模型越来越简单，CPU 上优化的边际效应越来越低。美团基于外部深度定制的 TensorFlow、NVIDIA HugeCTR，研发了 Booster GPU 训练架构。整体设计充分考虑算法、架构、新硬件的个性，并从数据、计算、通信等多个角度深度优化，比照之前 CPU 的工作，性价比晋升到 2~4 倍。从性能和齐备性上反对 TensorFlow 的各类训练接口（Train/Evaluate/Rredict 等)，反对 CPU 和 GPU 模型互相导入。易用性上 TensorFlow CPU 工作只须要一行代码就可实现 GPU 架构迁徙。目前在美团外卖举荐场景实现了大规模的投产利用，后续咱们将会全面推广到到家搜寻举荐技术部以及美团全业务线。

当然，Booster 基于 NVIDIA A100 单机多卡还有不少优化空间，如数据层面的样本压缩、序列化、特色解析，计算层面的多图算子调度、动静 shape 算子的编译优化，通信层面的量化通信等等。同时为了更宽泛的反对美团内的业务模型，Booster 的下一个版本也会反对更大的模型，以及多机多卡的 GPU 训练。

7 作者简介

家恒、国庆、峥少、晓光、鹏鹏、永宇、俊文、正阳、瑞东、翔宇、秀峰、王庆、封宇、事峰、黄军等，来自美团根底研发平台 - 机器学习平台训练引擎 & 到家研发平台 - 搜寻举荐技术部 Booster 联结项目组。

8 参考文献

• [1] https://tech.meituan.com/2021/12/09/meituan-tensorflow-in-recommender-systems.html
• [2] https://images.nvidia.cn/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf
• [3] https://www.usenix.org/system/files/conference/osdi14/osdi14-paper-li\_mu.pdf
• [4] https://github.com/NVIDIA-Merlin/HugeCTR
• [5] https://en.wikipedia.org/wiki/Nvidia\_Tesla
• [6] https://www.nvidia.com/en-us/data-center/dgx-a100
• [7] https://github.com/horovod/horovod
• [8] https://github.com/NVIDIA/nccl
• [9] https://www.tensorflow.org/xla
• [10] Yann LeCun, John S. Denker, and Sara A. Solla. Optimal brain damage. In NIPS, pp. 598–605. Morgan Kaufmann, 1989.
• [11] Kenji Suzuki, Isao Horiba, and Noboru Sugie. A simple neural network pruning algorithm with application to filter synthesis. Neural Process. Lett., 13(1):43–53, 2001.
• [12] Suraj Srinivas and R. Venkatesh Babu. Data-free parameter pruning for deep neural networks. In BMVC, pp. 31.1–31.12. BMVA Press, 2015.
• [13] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
• [14] Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. Compositional embeddings using complementary partitions for memory-efficient recommendation systems. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 165-175. 2020.
• [15] https://mp.weixin.qq.com/s/fOA\_u3TYeSwAeI6C9QW8Yw
• [16] Jianxun Lian, Xiaohuan Zhou, Fuzheng Zhang, Zhongxia Chen, Xing Xie, and Guangzhong Sun. 2018. xDeepFM: Combining Explicit and Implicit Feature Interactions for Recommender Systems. arXiv preprint arXiv:1803.05170 (2018).
• [17] Weiping Song, Chence Shi, Zhiping Xiao, Zhijian Duan, Yewen Xu, Ming Zhang, and Jian Tang. Autoint: Automatic feature interaction learning via self-attentive neural networks. In Proceedings of the 28th ACM International Conference on Information and Knowledge Management, pp. 1161-1170. 2019.
• [18] Guorui Zhou, Weijie Bian, Kailun Wu, Lejian Ren, Qi Pi, Yujing Zhang, Can Xiao et al. CAN: revisiting feature co-action for click-through rate prediction. arXiv preprint arXiv:2011.05625 (2020).
• [19] Chun-Hao Chang, Ladislav Rampasek, and Anna Goldenberg. Dropout feature ranking for deep learning models. arXiv preprint arXiv:1712.08645 (2017).
• [20] Xu Ma, Pengjie Wang, Hui Zhao, Shaoguo Liu, Chuhan Zhao, Wei Lin, Kuang-Chih Lee, Jian Xu, and Bo Zheng. Towards a Better Tradeoff between Effectiveness and Efficiency in Pre-Ranking: A Learnable Feature Selection based Approach. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 2036-2040. 2021.
• [21] Bencheng Yan, Pengjie Wang, Jinquan Liu, Wei Lin, Kuang-Chih Lee, Jian Xu, and Bo Zheng. Binary Code based Hash Embedding for Web-scale Applications. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 3563-3567. 2021.
• [22] Xiangyu Zhao, Haochen Liu, Hui Liu, Jiliang Tang, Weiwei Guo, Jun Shi, Sida Wang, Huiji Gao, and Bo Long. Autodim: Field-aware embedding dimension searchin recommender systems. In Proceedings of the Web Conference 2021, pp. 3015-3022. 2021.
• [23] Bencheng Yan, Pengjie Wang, Kai Zhang, Wei Lin, Kuang-Chih Lee, Jian Xu, and Bo Zheng. Learning Effective and Efficient Embedding via an Adaptively-Masked Twins-based Layer. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 3568-3572. 2021.
• [24] Ting Chen, Lala Li, and Yizhou Sun. Differentiable product quantization for end-to-end embedding compression. In International Conference on Machine Learning, pp. 1617-1626. PMLR, 2020.
• [25] Wang-Cheng Kang, Derek Zhiyuan Cheng, Ting Chen, Xinyang Yi, Dong Lin, Lichan Hong, and Ed H. Chi. Learning multi-granular quantized embeddings for large-vocab categorical features in recommender systems. In Companion Proceedings of the Web Conference 2020, pp. 562-566. 2020.
• [26] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836 (2016).
• [27] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. Advances in neural information processing systems 30 (2017).
• [28] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 (2017).
• [29] Chao Peng, Tete Xiao, Zeming Li, Yuning Jiang, Xiangyu Zhang, Kai Jia, Gang Yu, and Jian Sun. Megdet: A large mini-batch object detector. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pp. 6181-6189. 2018.

浏览美团技术团队更多技术文章合集

前端 | 算法 | 后端 | 数据 | 平安 | 运维 | iOS | Android | 测试

| 在公众号菜单栏对话框回复【2021 年货】、【2020 年货】、【2019 年货】、【2018 年货】、【2017 年货】等关键词，可查看美团技术团队历年技术文章合集。

| 本文系美团技术团队出品，著作权归属美团。欢送出于分享和交换等非商业目标转载或应用本文内容，敬请注明“内容转载自美团技术团队”。本文未经许可，不得进行商业性转载或者应用。任何商用行为，请发送邮件至 tech@meituan.com 申请受权。