关于美团:TensorFlow在推荐系统中的分布式训练优化实践

美团外部深度定制的 TensorFlow 版本，基于原生 TensorFlow 1.x 架构与接口，从大规模稠密参数的反对、训练模式、分布式通信优化、流水线优化、算子优化交融等多维度进行了深度优化。在举荐零碎场景，分布式扩展性晋升 10 倍以上，单位算力性能也有显著晋升，并在美团外部业务中大量应用，本文介绍了相干的优化与实际工作。

1 背景

TensorFlow（下文简称 TF）是谷歌推出的一个开源深度学习框架，在美团举荐零碎场景中失去了宽泛的应用。但 TensorFlow 官网版本对工业级场景的反对，目前做得并不是特地的欠缺。美团在大规模生产落地的过程中，遇到了以下几方面的挑战：

• 所有参数都是用 Variable 表白，对于百亿以上的稠密参数开拓了大量的内存，造成了资源的节约；
• 只反对百级别 Worker 的分布式扩大，对上千 Worker 的扩展性较差；
• 因为不反对大规模稠密参数动静增加、删除，增量导出，导致无奈反对 Online Learning；
• 大规模集群运行时，会遇到慢机和宕机；因为框架层不能解决，导会致工作运行异样。

以上这些问题，并不是 TensorFlow 设计的问题，更多是底层实现的问题。思考到美团大量业务的应用习惯以及社区的兼容性，咱们基于原生 TensorFlow 1.x 架构与接口，从大规模稠密参数的反对、训练模式、分布式通信优化、流水线优化、算子优化交融等多维度进行了深度定制，从而解决了该场景的外围痛点问题。

首先新零碎在反对能力层面，目前能够做到千亿参数模型，上千 Worker 分布式训练的近线性减速，全年样本数据可能 1 天内实现训练，并反对 Online Learning 的能力。同时，新零碎的各种架构和接口更加敌对，美团外部包含美团外卖、美团优选、美团搜寻、广告平台、公众点评 Feeds 等业务部门都在应用。本文将重点介绍大规模分布式训练优化的工作，心愿对大家可能有所帮忙或启发。

2 大规模训练优化挑战

2.1 业务迭代带来的挑战

随着美团业务的倒退，举荐零碎模型的规模和复杂度也在快速增长，具体表现如下：

• 训练数据：训练样本从到百亿增长到千亿，增长了近 10 倍。
• 稠密参数：个数从几百到几千，也增长了近 10 倍；总参数量从几亿增长到百亿，增长了 10~20 倍。
• 模型复杂度：越来越简单，模型单步计算工夫增长 10 倍以上。

对于大流量业务，一次训练试验，从几个小时增长到了几天，而此场景一次试验放弃在 1 天之内是根本的需要。

2.2 零碎负载剖析

2.2.1 问题剖析工具链

TensorFlow 是一个十分宏大的开源我的项目，代码有几百万行之多，原生零碎的监控指标太粗，且不反对全局的监控，如果要定位一些简单的性能瓶颈点，就比拟艰难。咱们基于美团曾经开源的监控零碎 CAT[2]，构建了 TensorFlow 的细粒度监控链路（如下图 1 所示），能够精准定位到性能的瓶颈问题。

同时，在性能优化的过程中，会波及到大量的性能测试和后果剖析，这也是一个十分消耗人力的工作。咱们形象了一套自动化的试验框架（如下图 2 所示），能够自动化、多轮次地进行试验，并主动采集各类监控指标，而后生成报告。

2.2.2 业务视角的负载剖析

在举荐零碎场景中，咱们应用了 TensorFlow Parameter Server[3]（简称 PS）异步训练模式来反对业务分布式训练需要。对于这套架构，上述的业务变动会带来什么样的负载变动？如下图 3 所示：

总结来看，次要包含通信压力、PS 并发压力、Worker 计算压力。对于分布式系统来说，通常是通过横向扩大来解决负载问题。尽管看来起能够解决问题，但从试验后果来看，当 PS 扩大到肯定数量后，单步训练工夫反而会减少，如下图 4 所示：

导致这种后果的外围起因是：Worker 单步训练须要和所有的 PS 通信同步实现，每减少 1 个 PS 要减少 N 条通信链路，这大大增加了链路提早（如下图 5 所示）。而一次训练要执行上百万、上千万步训练。最终导致 链路提早超过了加 PS 算力并发的收益。

而对于这个零碎，优化的外围难点在于：如何在无限的 PS 实例下，进行分布式计算的优化。

3 优化实际

3.1 大规模稠密参数介绍

对于举荐零碎模型，绝大多数参数都是稠密参数，而对稠密参数来说有一个十分重要的操作是 Embedding，这个操作通常也是负载最重的，也是后续优化的重点。因为咱们对稠密参数进行了从新定义，后续的优化也基于此之上，所以咱们先介绍一下这部分的工作。

在原生的 TensorFlow 中构建 Embedding 模块，用户须要首先创立一个足够装得下所有稠密参数的 Variable，而后在这个 Variable 上进行 Embedding 的学习。然而，应用 Variable 来进行 Embedding 训练存在很多弊病：

• Variable 的大小必须提前设定好，对于百亿千亿的场景，该设定会带来微小的空间节约；
• 训练速度慢，无奈针对稠密模型进行定制优化。

咱们首先解决了有无的问题，应用 HashTable 来代替 Variable，将稠密特色 ID 作为 Key，Embedding 向量作为 Value。相比原生应用 Variable 进行 Embedding 的形式，具备以下的劣势：

1. HashTable 的大小能够在训练过程中主动伸缩，防止了开拓冗余的存储空间，同时用户无需关注申请大小，从而升高了应用老本。
2. 针对 HashTable 计划施行了一系列定制优化，训练速度相比 Variable 有了很大的进步，能够进行千亿规模模型的训练，扩展性较好。
3. 得益于稠密参数的动静伸缩，咱们在此基础上反对了 Online Learning。
4. API 设计上放弃与社区版本兼容，在应用上简直与原生 Variable 统一，对接老本极低。

简化版的基于 PS 架构的实现示意如下图 6 所示：

外围流程大抵能够分为以下几步：

1. 稠密特色 ID（通常咱们会提前完成对立编码的工作）进入 Embedding 模块，借助 TensorFlow 搭建的 Send-Recv 机制，这些稠密特色 ID 被拉取到 PS 端，PS 端上的 Lookup 等算子会理论从底层 HashTable 中查问并组装 Embedding 向量。
2. 上述 Embedding 向量被 Worker 拉回进行后续训练，并通过反向流传计算出这部分参数的梯度，这些梯度进一步被位于 PS 端的优化器拉回。
3. PS 端的优化器首先调用 Find 算子，从 HashTable 获取到梯度对应的原始稠密参数向量和相应的优化器参数，最终通过优化算法，实现对 Embedding 向量和优化器参数的更新计算，再通过 Insert 算子插入 HashTable 中。

3.2 分布式负载平衡优化

这部分优化，是分布式计算的经典优化方向。PS 架构是一个典型的“水桶模型”，为了实现一步训练，Worker 端须要和所有 PS 实现交互，因而 PS 之间的均衡就显得十分重要。然而在实践中，咱们发现多个 PS 的耗时并不平衡，其中的起因，既包含 TensorFlow PS 架构简略的切图逻辑（Round-Robin）带来的负载不平衡，也有异构机器导致的不平衡。

对于举荐模型来说，咱们的次要优化策略是，把所有稠密参数和大的浓密参数主动、平均的切分到每个 PS 上，能够解决大多数这类问题。而在实际过程中，咱们也发现一个比拟难排查的问题：原生 Adam 优化器，实现导致 PS 负载不平衡。上面会具体介绍一下。

在 Adam 优化器中，它的参数优化过程须要两个 β 参加计算，在原生 TensorFlow 的实现中，这两个 β 是所有须要此优化器进行优化的 Variabl（或 HashTable）所共享的，并且会与第一个 Variable（名字字典序）落在同一个 PS 下面，这会带来一个问题：每个优化器只领有一个 β_1 和一个 β_2，且仅位于某个 PS 上。因而，在参数优化的过程中，该 PS 会接受远高于其余 PS 的申请，从而导致该 PS 成为性能瓶颈。

然而通过观察 Adam 的优化算法，咱们能够看到 β_1 和 β_2 都是常量，且蓝色高亮的局部都是绝对独立的计算过程，各个 PS 之间能够独立实现。基于这样的发现，优化的办法也就十分直观了，咱们为每一个 PS 上的 Adam 优化器冗余创立了 β 参数，并在本地计算 t 和 alpha 值，去除了因而负载不均导致的 PS 热点问题。

该优化所带来的晋升具备普适性且成果显著，在美团外部某业务模型上，通过 β 热点去除能够带来 9% 左右的性能晋升。此外，因为解脱了对 β 的全局依赖，该优化还能进步 PS 架构的可扩展性，在扩增 Worker 数量的时候相比之前会带来更好的减速比。

3.3 通信优化

通过 2.2 章节的剖析可知，零碎的通信压力也十分大，咱们次要基于 RDMA 做了通信优化的工作。首先简略介绍一下 RDMA，相比拟于传统基于套接字 TCP/IP 协定栈的通信过程，RDMA 具备零拷贝、内核旁路的劣势，不仅升高了网络的提早，同时也升高了 CPU 的占用率，RDMA 更适宜深度学习模型的相干通信过程。

RDMA 次要包含三种协定 Infiniband、RoCE(V1, V2)、iWARP。在美团外部的深度学习场景中，RDMA 通信协议应用的是 RoCE V2 协定。目前在深度学习训练畛域，尤其是在浓密模型训练场景（NLP、CV 等），RDMA 曾经是大规模分布式训练的标配。然而，在大规模稠密模型的训练中，开源零碎对于 RDMA 的反对十分无限，TensorFlow Verbs[4]通信模块曾经很长时间没有更新了，通信成果也并不现实，咱们基于此之上进行了很多的改良工作。

通过优化后的版本，在 1TB Click Logs[5]公开数据集、DLRM[6]模型、100 个 Worker 以上的训练，性能晋升了 20%~40%。在美团的多个业务模型上，比照 TensorFlow Seastar[7]革新的通信层实现也有 10%~60% 的速度晋升。同时也把咱们的工作回馈给了社区。

3.3.1 Memory Registration 优化

RDMA 有三种数据传输的形式 SEND/RECV、WRITE、READ，其中 WRITE、READ 相似于数据发送方间接在近程 Memory 进行读写，Receiver 无奈感知，WRITE 和 READ 实用于批量数据传输。在 TensorFlow 外部，基于 RDMA 的数据传输方式应用的是 WRITE 单边通信模式。

在 RDMA 传输数据时，须要提前开拓内存空间并将其注册到网卡设施上（Memory Registration 过程，下称 MR），使得这片空间能够被网卡间接操作。开拓新的内存并注册到设施上，整个过程是比拟耗时的。下图 9 展现了不同大小的内存绑定到网卡设施上的耗时，能够看到随着注册内存的增大，绑定 MR 的耗时迅速减少。

社区版 Tensorflow RDMA 实现，Tensor 创立仍旧沿用了对立的 BFC Allocator，并将所有创立的 Tensor 都注册到 MR 上。正如下面所提到的，MR 的注册绑定具备性能开销，高频、大空间的 MR 注册会带来显著的性能降落。而训练过程中的 Tensor，只有那些波及到跨节点通信的 Tensor 有必要进行 MR，其余 Tensor 并不需要注册到 MR。因而，优化的办法也就比拟间接了，咱们辨认并治理那些通信 Tensor，仅对这些跨节点通信的 Tensor 进行 MR 注册就好了。

3.3.2 RDMA 动态分配器

RDMA 动态分配器是上一个 MR 注册优化的延长。通过 Memory Registration 优化，去除非传输 Tensor 的 MR 注册，咱们升高了 MR 注册数量。然而在稠密场景大规模的训练下，并行训练的 Worker 常有几百上千个，这会带来新的问题：

• PS 架构中的 PS 和 Worker 互为 Client-Server，这里以 PS 端为例，当 Worker 数目减少到上千个时，Worker 数目的增多，造成 PS 端 MR 注册频次还是十分高，减少了内存调配注册的耗时。
• 因为稠密场景不同 Step 之间同一个算子输入 Tensor 的形态可能发生变化，导致了创立的 MR 可复用性较差，带来了较高的内存碎片和反复注册 MR 开销。

针对下面的问题，咱们引入了 MR 动态分配器的策略。

这里外围的设计思路为：

1. 尽管稠密场景同一个算子输入 Tensor 的 Shape 存在变动的可能，然而整体变动幅度可控，通过监控与剖析，是能够找到一个较为稳固的内存大小，满足多 Step 间 Tensor 的存储需要。
2. 基于下面的信息，咱们批改了原有逐 Tensor(Request)的 MR 申请策略，通过一次性预申请一块较大的空间并注册到网卡端，后续通过本人保护的调配策略进行空间的调配，大大降低了 MR 申请的频率，绝大多数状况下，训练全过程中只须要一次 MR 注册申请即可。
3. 咱们引入了一种简略的替换协定，将传输 Tensor 的 Shape，Data 打包到一起，写到 Client 端。Client 端依据协定，解析出 Tensor 大小，并最终读取 Data，防止了原生实现中因 Tensor 的 Shape 变动而产生的屡次协商过程。

具体到实现中，咱们引入了 Allocation Analysis 模块，在训练开始的一段时间，咱们会对调配的历史数据进行剖析，以失去一个理论预开拓 MR 大小以及各个 Tensor 的预留空间大小。而后咱们会暂停训练的过程，启动 Allocator 的结构过程，包含 MR 的创立以及通信双端的信息同步。利用相干信息结构 MR Info Map，这个 Map 的 Key 是传输 Tensor 的惟一标记（ParsedKey，计算图切图时确定），Info 构造体中蕴含了本地地址指针、offset 大小、ibv_send_wr 相干信息等。而后复原训练，后续 Tensor 的传输就能够应用动态开拓好的 MR 进行收发，也免去了因 Shape 变动而产生的屡次协商过程。

3.3.3 Multi RequestBuffer 与 CQ 负载平衡

TensorFlow 社区版的 RDMA 通信过程，不仅仅蕴含下面 Tensor 数据的发送和接管过程，还包含传输相干的管制音讯的发送和接管过程，管制音讯的发送和接管过程同样是应用了 ibv_post_send 和 ibv_post_recv 原语。原生的控制流实现存在一些瓶颈，在大规模训练时会限度控制流的吞吐，进而影响数据收发的效率。具体体现在：

• 申请的发送通过同一片 RequestBuffer 内存进行写出，多个 Client 的申请均依赖这一片 Buffer，也就导致到控制流信息理论是串行发送的，只有等到对端的 Ack 信息后，才能够下一个 Request 的写出，限度了申请的发送吞吐。
• 在 Client 端须要轮询 RDMA Completion Queue 来取得申请的达到，以及相干状态的变更。原生实现仅有一个 Completion Queue，单线程进行轮询解决，在大规模分布式训练中，限度了应答的效率。

针对下面的问题，咱们采纳了 Multi RequestBuffer 与 CQ 负载平衡优化，破除了在申请发送和申请应答环节可能存在的吞吐瓶颈。

3.3.4 Send-Driven & Rendezvous-Bypass

对于 Tensorflow PS 架构相熟的同学会理解，一整张计算图被切割为 Worker 端和 PS 端后，为了使两张计算图可能彼此替换数据，建设了基于 Rendezvous（汇合点）机制的异步数据交换模式。如下图 12 所示：

基于上图的切图逻辑，Recv 算子代表着这一侧计算图有 Tensor 的需要，而 Tensor 的生产者则位于与之配对的另一设施上的 Send 算子背地。

在具体实现上，Tensorflow 实现了 Recv-Driven 的数据交换模式，如上图所示，位于 DeviceA 和 DeviceB 的两张计算图会异步并发的执行，位于 DeviceB 的 Recv 执行时会发动一条 RPC 申请发往 DeviceA，DeviceA 收到申请后，会将申请路由到 Rendezvous 中，如果在当中发现所须要的数据曾经生产好，并被 Send 算子注册了进来，那么就地获取数据，返回给 DeviceB；如果此时数据还没有生产好，则将来自于 DeviceB 的 Recv 申请注册在 Rendezvous 中，期待后续 DeviceA 生产好后，由 Send 算子发送过去，找到注册的 Recv，触发回调，返回数据给 DeviceB。

咱们看到，汇合点机制优雅地解决了生产者消费者节奏不同状况下数据交换的问题。不过 Recv-Driven 的模式也引入了两个潜在的问题：

• 据咱们的察看，在理论业务模型中，在 Rendezvous 中 Recv 算子期待 Send 算子的比例和 Send 算子期待 Recv 算子的比例相当，也就是说对于 Send 等到 Recv 的数据，在 Send 筹备好的那一刹那就能够发给对端，然而因为机制实现问题，还是期待 Recv 算子过去，才将数据拉取回去，通信过程耗时较长。
• Rendezvous 作为一个数据交换的热点，它外部的逻辑开销并不低。

针对下面提到的问题，咱们在 RDMA 上实现了另外一种数据交换的模式，叫做 Send-Driven 模式。与 Recv-Driven 模式绝对，顾名思义就是有 Send 算子间接将数据写到 Recv 端，Recv 端接收数据并注册到本地 Rendezvous 中，Recv 算子间接从本地的 Rendezvous 中获取数据。具体流程如下图 13 所示：

从图中能够看到，相较于 Recv-Driven 模式，Send-Driven 模式的通信流程失去了比拟大的简化，另外在数据 ready 后立刻发送的个性，跳过了一侧的 Rendezvous，并且对于生产者先于消费者的状况，能够放慢生产端数据获取的速度。

3.4 提早优化

这部分优化，也是分布式计算的经典优化方向。整个流程链路上那些能够精简、合并、重叠须要一直去开掘。对于机器学习零碎来说，相比其它的零碎，还能够用一些近似的算法来做这部分工作，从而取得较大的性能晋升。上面介绍咱们在两个这方面做的一些优化实际。

3.4.1 稠密域参数聚合

在启用 HashTable 存储稠密参数后，对应的，一些配套参数也须要替换为 HashTable 实现，这样整个计算图中会呈现多张 HashTable 以及大量的相干算子。在实践中，咱们发现须要尽量升高 Lookup/Insert 等算子的个数，一方面升高 PS 的负载，一方面升高 RPC QPS。因而，针对稠密模型的常见用法，咱们进行了相干的聚合工作。

以 Adam 优化器为例，须要创立 m、v 两个 slot，以保留优化中的动量信息，它的 Shape 与 Embedding 雷同。在原生优化器中，这两个 Variable 是独自创立的，并在反向梯度更新的时候会去读写。同理，应用 HashTable 计划时，咱们须要同时创立两张独自的 HashTable 用来训练 m、v 参数。那么在前向，反向中须要别离对 Embedding、m、v 进行一次 Lookup 和一次 Insert，总共须要三次 Lookup 和三次 Insert。

这里一个优化点就是将 Embedding、m、v，以及低频过滤的计数器（见下图 14 的 Counting HashTable）聚合到一起，作为 HashTable 的 Value，这样对稠密参数的相干操作就能够聚合执行，大大减少了稠密参数操作频次，升高了 PS 的压力。

该个性属于一个普适型优化，开启聚合性能后，训练速度有了显著的进步，性能晋升幅度随着模型和 Worker 规模的变动，成果总是正向的。在美团外部实在业务模型上，聚合之后性能相比非聚合形式能晋升了 45% 左右。

3.4.2 Embedding 流水线优化

流水线，在工业生产中，指每一个生产单位只专一解决某个片段的工作，以进步工作效率及产量的一种生产方式。在计算机领域内，更为大家熟知的是，流水线代表一种多任务之间 Overlap 执行的并行化技术。例如在典型的 RISC 处理器中，用户的程序由大量指令形成，而一条指令的执行又能够大抵分为：取指、译码、执行、访存、写回等环节。这些环节会利用到指令 Cache、数据 Cache、寄存器、ALU 等多种不同的硬件单元，在每一个指令周期内，这 5 个环节的硬件单元会并行执行，得以更加充沛的利用硬件能力，以此进步整个处理器的指令吞吐性能。处理器的指令流水线是一套简单而零碎的底层技术，但其中的思维在分布式深度学习框架中也被大量的应用，例如：

• 如果将分布式训练简略的形象为计算和通信两个过程，绝大多数支流的深度学习框架都反对在执行计算图 DAG 时，通信和计算的 Overlap。
• 如果将深度模型训练简略的分为前向和反向，在单步内，因为两者的强依赖性，无奈做到无效并行，字节 BytePS[8]中引入的通信调度突破了 step iteration 间的屏障，上一轮的局部参数更新结束后，即可提前开始下轮的前向计算，加强了整体视角下前反向的 Overlap。
• 百度 AIBox[9]为了解决 CTR 场景 GPU 训练时，参数位于主存，但计算位于 GPU 的问题，奇妙调度不同硬件设施，搭建起了次要利用 CPU/ 主存 / 网卡的参数预筹备阶段和次要利用 GPU/NVLink 的网络计算阶段，通过两个阶段的 Overlap 达到更高的训练吞吐。

咱们看到，在深度学习框架设计上，通过剖析场景，能够从不同的视角挖掘可并行的阶段，来进步整体的训练吞吐。

对于大规模稠密模型训练时，外围模型流程是：先执行稠密参数的 Embedding，而后执行浓密局部子网络。其中稠密参数 Embedding 在远端 PS 上执行，次要消耗网络资源，而浓密局部子网络在本地 Worker 执行，次要消耗计算资源。这两局部占了整个流程的大部分工夫，在美团某理论业务模型上别离耗时占比：40+%、50+%。

那咱们是否能够提前执行稠密参数的 Embedding，来做到通信和计算的 Overlap，暗藏掉这部分工夫呢？从零碎实现上必定是可行的，但从算法上讲，这样做会引入参数 Staleness 的问题，可能会导致模型精度受到影响。但在理论的生产场景中，大规模异步训练时自身就会带来几十到几百个步的滞后性问题。通过咱们测试，提前获取一两步的稠密参数，模型精度并未受到影响。

在具体实现上，咱们把整个计算图拆分为 Embedding Graph（EG）和 Main Graph（MG）两张子图，两者异步独立执行，做到拆分流程的 Overlap（整个拆分过程，能够做到对用户通明）。EG 次要笼罩从样本中抽取 Embedding Key，查问组装 Embedding 向量，Embedding 向量更新等环节。MG 次要蕴含浓密局部子网络计算、梯度计算、浓密参数局部更新等环节。

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

上图是 Embedding 流水线的架构流程图。直观来看分为左侧的样本散发模块，顶部的跨 Session 数据交换模块，以及主动图切分失去的 Embedding Graph 和 Main Graph，蓝色的圆圈代表新增算子，橙色箭头代表 EG 重点流程，蓝色箭头代表 MG 重点流程，红色箭头代表样本数据重点流程。

1. 以对用户通明的模式引入了一层名为 Pipeline Dataset 的形象层，这一层的产生是为了满足 EG/MG 两张计算图以不同节奏运行的需要，反对自定义配置。另外，为了使得整个流水线中的数据做到彼此的配套，这里还会负责进行一个全局 Batch ID 的生成及注册工作。Pipeline Dataset 对外裸露两种 Iterator，一个供 EG 应用，一个供 MG 应用。Pipeline Dataset 底部共享 TensorFlow 原生的各层 Dataset。
2. 顶部的 ExchangeManager 是一个动态的，跨 Session 的数据交换媒介，对外裸露数据注册和数据拉取的能力。形象这个模块的起因是，EG 和 MG 本来归属于一张计算图，因为流水线的起因拆解为拆为两张图，这样咱们须要建设一种跨 Session 的数据交换机制，并精确进行配套。它外部以全局 Batch ID 做 Key，前面治理了样本数据、Embeding 向量、Embedding 梯度、Unique 后的 Index 等数据，并负责这些数据的生命周期治理。
3. 两头的 Embedding Graph 由独立的 TF Session 运行于一个独立的线程中，通过 a 算子取得样本数据后，进行特色 ID 的抽取等动作，并进行基于 HashTable 办法的稠密参数查问，查问后果通过 c 算子搁置到 ExchangeManager 中。EG 中还蕴含用于反向更新的 f 算子，它会从 ExchangeManager 中获取 Embedding 梯度和与其配套的前向参数，而后执行梯度更新参数逻辑。
4. 上面的 Main Graph 负责理论浓密子网络的计算，咱们继承并实现一种可训练的 EmbeddingVariable，它的构建过程（d 算子）会从 ExchangeManager 查找与本人配套的 Embedding 向量封装成 EmbeddingVariable，给浓密子网络。此外，在 EmbeddingVariable 注册的反向办法中，咱们增加了 e 算子使得 Embedding 梯度得以增加到 ExchangeManager 中，供 EG 中的 f 算子生产。

通过下面的设计，咱们就搭建起了一套可控的 EG/MG 并发流水线训练模式。总体来看，Embedding 流水线训练模式的收益起源有：

• 通过咱们对多个业务模型的 Profiling 剖析发现，EG 和 MG 在工夫的比例上在 3:7 或 4:6 的左右，通过将这两个阶段并行起来，能够无效的暗藏 Embedding 阶段，使得 MG 网络计算局部简直总是能够立刻开始，大大减速了整体模型的训练吞吐。
• TensorFlow 引擎中当应用多个优化器（稠密与非稠密）的时候，会呈现反复构建反向计算图的问题，肯定水平减少了额定计算，通过两张子图的拆分，恰好防止了这个问题。
• 在施行过程中的 ExchangeManager 不仅负责了 Embedding 参数和梯度的替换，还承当了元数据复用治理的职责。例如 Unique 等算子的后果保留，进一步升高了反复计算。

另外，在 API 设计上，咱们做到了对用户通明，仅需一行代码即可开启 Embedding 流水线性能，对用户暗藏了 EG/MG 的切割过程。目前，在美团某业务训练中，Embedding 流水线性能在 CPU PS 架构下能够带来 20%~60% 的性能晋升（而且 Worker 并发规模越大，性能越好）。

3.5 单实例 PS 并发优化

通过 2.2 章节的剖析可知，咱们不能通过继续扩 PS 来晋升分布式工作的吞吐，单实例 PS 的并发优化，也是十分重要的优化方向。咱们次要的优化工作如下。

3.5.1 高性能的 HashTable

PS 架构下，大规模稠密模型训练对于 HashTable 的并发读写要求很高，因为每个 PS 都要承当成百乃至上千个 Worker 的 Embedding 压力，这里咱们综合速度和稳定性思考，选用了 tbb::concurrent_hash_map[10]作为底层 HashTable 表实现，并将其包装成一个新的 TBBConcurrentHashTable 算子。通过测试，在千亿规模下 TBBConcurrentHashTable 比原生 MutableDenseHashTable 训练速度上快了 3 倍。

3.5.2 HashTable BucketPool

对于大规模稠密模型训练来说，Embedding HashTable 会面对大量的并发操作，通过 Profiling 咱们发现，频繁动静的内存申请会带来了较大性能开销（即便 TensorFlow 的 Tensor 有专门的内存分配器）。咱们基于内存池化的思路优化了 HashTable 的内存治理。

咱们在 HashTable 初始化时，会先为 Key 和 Value 别离发明两个 BucketPool，每个池子都会先 Malloc 较大一块内存备用，思考到可能会有对 HashTable 进行中的 Key 和 Value 进行 Remove 的场景（如 Online Learning 训练时），须要对从 HashTable 中删除的 Key 和 Value 所应用的内存进行回收，因而每个 BucketPool 还有一个 ReuseQueue 来负责保护回收的内存。每次向外部的哈希表数据结构中 Insert Key 和 Value 的时候，Key 和 Value 内存和开释调配都进行池化治理。用这种形式升高了大规模稠密训练中遇到稠密内存调配开销，整体端到端训练性能晋升了 5% 左右。

3.6 单位算力吞吐优化

通过 2.2 章节的剖析，Worker 的计算压力也十分大，如果不优化 Worker，同时要放弃吞吐，须要横向扩大更多的 Worker，给 PS 带来更大的压力。而对于用户来说，如果能在无限的计算资源下带来性能晋升，对业务价值更高。咱们通过 CAT 统计出了一些高频算子，并进行了专项优化。这里选取 Unique&DynamicPartition 算子交融案例进行分享。

在 TensorFlow PS 架构中，包含 Embedding 向量在内的共享参数都存储在 PS 上，并通过网络与 Worker 交互，在进行 Embedding 查问过程中，往往会波及如下两个环节：

• 因为稠密参数的性质，从样本中抽取失去的待查问 Embedding ID，它的反复率往往高达 70%~90%，如果不进行去重查问，不论是对 HashTable 的查问还是网络的传输，都会带来不小的压力。因而，通常会在查问前进行 Unique 操作。
• 在大规模稠密场景中，为了存储千亿规模的参数，会有多个 PS 机器独特承载。而 Worker 端会负责对查问申请依照设定的路由规定进行切分，这里通常会在查问前进行 DynamicPartition 动作。

通常这两个过程会利用 TensorFlow 既有的算子进行搭建，但在理论应用中，咱们发现它并不是很高效，次要问题在于：

• Unique 算子原生实现，它外部应用的内存调配策略较为低效。应用了两倍输出参数（Embedding ID）的大小进行内存调配，但因为输出参数较大，而且反复率高，导致 HashTable 创立过大且十分稠密。简直每次插入都会产生一次 minor_page_fault，导致 HashTable 性能降落。咱们应用 Intel Vtune 验证了这一点（参见图 18）。
• Unique 和 Dynamic Partition 算子存在冗余数据遍历，这些操作其实能够在一次数据遍历中全副做完，节俭掉算子切换、冗余数据遍历的耗时。

总结来说，HashTable 开拓过大会导致大量的 minor_page_fault，导致访存的工夫减少，HashTable 过小又可能会导致扩容。咱们采纳了基于启发式算法的内存自适应 Unique 算子实现，通过对训练历史反复率的统计，咱们能够失去一个绝对正当的 HashTable 大小，来进步访存的性能；另外 Unique 算子内 HashTable 的具体抉择上，通过咱们的多种测试，抉择了 Robin HashTable 替换了原生 TF 中的实现。

进一步，咱们对围绕 Embedding ID 的 Unique 和 Partition 环节进行了算子合并，简化了逻辑实现。通过上述的优化，Unique 单算子能够获得 51% 的减速，在实在模型端到端上能够取得 10% 左右的性能晋升，算子总数量升高了 4%。

在整个要害算子优化的过程中，Intel 公司的林立凡、张向泽、高超进行大量的技术支持，咱们也复用了他们的局部优化工作，在此深表感激！

4 大规模稠密算法建模

大规模稠密能力在业务落地的过程中，算法层面还须要从特色和模型构造上进行对应降级，能力拿到十分好的成果。其中外卖广告从业务特点登程，引入大规模稠密特色实现外卖场景下特色体系的降级，提供了更高维的特色空间和参数空间，加强了模型的拟合能力。从新设计了面向高维稠密场景的特色编码方案，解决了特色编码过程中的特色抵触问题，同时编码过程去掉了局部冗余的特色哈希操作，肯定水平上简化了特色解决逻辑，并升高了特色计算的耗时。

在零碎层面，面对百亿参数、百亿样本以上量级的大规模稠密模型的训练，会带来训练迭代效率的大大降低，单次试验从一天以内，增长到一周左右。美团机器学习平台训练引擎团队，除了上述 TensorFlow 框架层面的优化、还针对业务模型进行了专项优化，整体吞吐优化了 8 到 10 倍（如果投入更多计算资源，能够进一步减速），大大晋升业务的迭代效率，助力外卖广告业务获得了较为显著的晋升。

5 总结与瞻望

TensorFlow 在大规模举荐零碎中被宽泛应用，但因为不足大规模稠密的大规模分布式训练能力，妨碍了业务的倒退。美团基于 TensorFlow 原生架构，反对了大规模稠密能力，并从多个角度进行了深度优化，做到千亿参数、千亿样本高效的分布式训练，并在美团外部进行了大规模的应用。对于这类要害能力的缺失，TensorFlow 社区也引起了共鸣，社区官网在 2020 年创立了 SIG Recommenders[11]，通过社区共建的形式来解决此类问题，美团后续也会踊跃的参加到社区的奉献当中去。

美团举荐零碎场景的模型训练，目前次要运行在 CPU 上，但随着业务的倒退，有些模型变得越来越简单，CPU 上曾经很难有优化空间（优化后的 Worker CPU 使用率在 90% 以上）。而近几年，GPU 的计算能力突飞猛进，新一代的 NVIDIA A100 GPU，算力达到了 156TFLOPS（TF32 Tensor Cores）、80G 显存、卡间带宽 600GB/s。对于这类简单模型的 Workload，咱们基于 A100 GPU 架构，设计了下一代的分布式训练架构，通过初步优化，在美团某大流量业务举荐模型上也拿到了较好的成果，目前还在进一步优化当中，后续咱们会进行分享，敬请期待。

6 作者简介

• 逸帆、家恒、峥少、鹏鹏、永宇、正阳、黄军等，来自美团根底研发平台，机器学习平台训练引擎组，次要负责美团分布式机器学习训练零碎的性能优化与能力建设。
• 海涛，来自美团外卖广告策略团队，次要负责美团外卖广告业务的算法摸索和策略落地工作。

7 招聘信息

美团机器学习平台大量岗位继续招聘中，社招 / 校招均可（欢送投递咱们的校招北斗岗位：美团机器学习平台基础架构），坐标北京 / 上海，构建多畛域的公司级机器学习平台，帮大家吃得更好，生存更好。简历可投递至：huangjun03@meituan.com。

8 参考信息

• [1] https://www.usenix.org/system/files/conference/osdi16/osdi16-abadi.pdf
• [2] https://github.com/dianping/cat
• [3] https://www.usenix.org/system/files/conference/osdi14/osdi14-paper-li_mu.pdf
• [4] https://github.com/tensorflow/networking/tree/master/tensorflow_networking/verbs
• [5] https://labs.criteo.com/2013/12/download-terabyte-click-logs/
• [6] https://arxiv.org/abs/1906.00091
• [7] https://github.com/tensorflow/networking/tree/master/tensorflow_networking/seastar
• [8] https://github.com/bytedance/byteps
• [9] http://research.baidu.com/Public/uploads/5e18a1017a7a0.pdf
• [10] https://github.com/oneapi-src/oneTBB
• [11] https://github.com/tensorflow/recommenders-addons

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

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

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

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