TensorFlow 在推荐系统中的分布式训练优化实践

[tfwechat (TensorFlow_WeChat)]

发布人：美团技术团队

美团内部深度定制的 TensorFlow 版本，基于原生 TensorFlow 1.x 架构与接口，从大规模稀疏参数的支持、训练模式、分布式通信优化、流水线优化、算子优化融合等多维度进行了深度优化。在推荐系统场景中，分布式扩展性提升 10 倍以上，单位算力性能也有显著提升，并在美团内部业务中大量使用，本文介绍了相关的优化与实践工作。

• 1 背景
• 2 业务迭代带来的挑战
  • 2.1 业务迭代带来的挑战
  • 2.2 系统负载分析
• 3 优化实践
  • 3.1 大规模稀疏参数介绍
  • 3.2 分布式负载均衡优化
  • 3.3 通信优化
  • 3.4 延迟优化
  • 3.5 单实例 PS 并发优化
  • 3.6 单位算力吞吐优化
• 4 大规模稀疏算法建模
• 5 总结与展望

背景

TensorFlow（下文简称 TF）是谷歌推出的一个开源深度学习框架，在美团推荐系统场景中得到了广泛的使用。但 TensorFlow 官方版本对工业级场景的支持，目前做得并不是特别的完善。美团在大规模生产落地的过程中，遇到了以下几方面的挑战：

• 所有参数都是用 Variable 表达， 对于百亿以上的稀疏参数开辟了大量的内存，造成了资源的浪费；

• 只支持百级别 Worker 的分布式扩展，对上千 Worker 的扩展性较差；

• 由于不支持大规模稀疏参数动态添加、删除，增量导出，导致无法支持 Online Learning；

• 大规模集群运行时，会遇到慢机和宕机；由于框架层不能处理，导会致任务运行异常。

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

首先新系统在支持能力层面，目前可以做到千亿参数模型，上千 Worker 分布式训练的近线性加速，全年样本数据能够1天内完成训练，并支持 Online Learning 的能力。同时，新系统的各种架构和接口更加友好，美团内部包括美团外卖、美团优选、美团搜索、广告平台、大众点评 Feeds 等业务部门都在使用。本文将重点介绍大规模分布式训练优化的工作，希望对大家能够有所帮助或启发。

大规模训练优化挑战

2.1 业务迭代带来的挑战

随着美团业务的发展，推荐系统模型的规模和复杂度也在快速增长，具体表现如下：

• 训练数据： 训练样本从到百亿增长到千亿，增长了近 10 倍。

• 稀疏参数： 个数从几百到几千，也增长了近 10 倍；总参数量从几亿增长到百亿，增长了 10~20 倍。

• 模型复杂度： 越来越复杂，模型单步计算时间增长 10 倍以上。

对于大流量业务，一次训练实验，从几个小时增长到了几天，而此场景一次实验保持在 1 天之内是基本的需求。

2.2 系统负载分析

2.2.1 问题分析工具链

TensorFlow 是一个非常庞大的开源项目，代码有几百万行之多，原生系统的监控指标太粗，且不支持全局的监控，如果要定位一些复杂的性能瓶颈点，就比较困难。我们基于美团已经开源的监控系统 CAT，构建了 TensorFlow 的细粒度监控链路（如下图 1 所示），可以精准定位到性能的瓶颈问题。

图 1 TensorFlow PS 架构全链路监控

同时，在性能优化的过程中，会涉及到大量的性能测试和结果分析，这也是一个非常耗费人力的工作。我们抽象了一套自动化的实验框架（如下图 2 所示），可以自动化、多轮次地进行实验，并自动采集各类监控指标，然后生成报告。

图 2 自动化实验框架

2.2.2 业务视角的负载分析

在推荐系统场景中，我们使用了 TensorFlow Parameter Server（简称 PS）异步训练模式来支持业务分布式训练需求。对于这套架构，上述的业务变化会带来什么样的负载变化？如下图 3 所示：

图 3 TensorFlow PS 架构大规模训练负载分析

总结来看，主要包括通信压力、PS 并发压力、Worker 计算压力。对于分布式系统来说，通常是通过横向扩展来解决负载问题。虽然看来起可以解决问题，但从实验结果来看，当 PS 扩展到一定数量后，单步训练时间反而会增加，如下图 4 所示：

图 4 扩展 PS 提升训练性能实验

导致这种结果的核心原因是：Worker 单步训练需要和所有的 PS 通信同步完成，每增加 1 个 PS 要增加 N 条通信链路，这大大增加了链路延迟（如下图 5 所示）。而一次训练要执行上百万、上千万步训练。最终导致链路延迟超过了加 PS 算力并发的收益。

图 5 增加 PS 带来的链路开销

而对于这个系统，优化的核心难点在于：如何在有限的 PS 实例下，进行分布式计算的优化。

优化实践

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 所示：

图 6 支撑大规模稀疏参数的 HashTable 方案

核心流程大致可以分为以下几步：

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 成为性能瓶颈。

图 7 Adam 优化算法

但是通过观察 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 通信模块已经很长时间没有更新了，通信效果也并不理想，我们基于此之上进行了很多的改进工作。

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

3.3.1 Memory Registration 优化

RDMA 有三种数据传输的方式 SEND/RECV、WRITE、READ，其中 WRITE、READ 类似于数据发送方直接在远程 Memory 进行读写，Receiver 无法感知，WRITE 和 READ 适用于批量数据传输。在 TensorFlow 内部，基于 RDMA 的数据传输方式使用的是 WRITE 单边通信模式。

图 8 RDMA 传输方式

在 RDMA 传输数据时，需要提前开辟内存空间并将其注册到网卡设备上（Memory Registration 过程，下称 MR），使得这片空间可以被网卡直接操作。开辟新的内存并注册到设备上，整个过程是比较耗时的。下图 9 展示了不同大小的内存绑定到网卡设备上的耗时，可以看到随着注册内存的增大，绑定 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 静态分配器的策略。

图 10 MR 静态分配器

这里核心的设计思路为：

1. 虽然稀疏场景同一个算子输出 Tensor 的 Shape 存在变化的可能，但是整体变化幅度可控，通过监控与分析，是可以找到一个较为稳定的内存大小，满足多 Step 间 Tensor 的存储需求。

2. 基于上面的信息，我们修改了原有逐 Tensor(Request) 的 MR 申请策略，通过一次性预申请一块较大的空间并注册到网卡端，后续通过自己维护的分配策略进行空间的分配，大大降低了 MR 申请的频率，绝大多数情况下，训练全过程中只需要一次 MR 注册申请即可。

3. 我们引入了一种简单的交换协议，将传输 Tensor 的 Shape，Data 打包到一起，写到 Client 端。Client 端根据协议，解析出 Tensor 大小，并最终读取 Data，避免了原生实现中因 Tensor 的 Shape 变化而产生的多次协商过程。

图 11 MR 静态分配器构造流程

具体到实现中，我们引入了 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 所示：

图 12 TensoFlow 切图之 Send-Recv 对添加

基于上图的切图逻辑，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 所示：

图 13 原生的 Recv-Driven 与补充的 Send-Driven 机制

从图中可以看到，相较于 Recv-Driven 模式，Send-Driven 模式的通信流程得到了比较大的简化，另外在数据 ready 后立即发送的特性，跳过了一侧的 Rendezvous，并且对于生产者先于消费者的情况，可以加快消费端数据获取的速度。

3.4 延迟优化

这部分优化，也是分布式计算的经典优化方向。整个流程链路上那些可以精简、合并、重叠需要不断去挖掘。对于机器学习系统来说，相比其它的系统，还可以用一些近似的算法来做这部分工作，从而获得较大的性能提升。下面介绍我们在两个这方面做的一些优化实践。

3.4.1 稀疏域参数聚合

在启用 HashTable 存储稀疏参数后，对应的，一些配套参数也需要替换为 HashTable 实现，这样整个计算图中会出现多张 HashTable 以及大量的相关算子。在实践中，我们发现需要尽量降低 Lookup/Insert 等算子的个数，一方面降低 PS 的负载，一方面降低 RPC QPS。因此，针对稀疏模型的常见用法，我们进行了相关的聚合工作。

以 Adam 优化器为例，需要创建两个 slot，以保存优化中的动量信息，它的 Shape 与 Embedding 相同。在原生优化器中，这两个 Variable 是单独创建的，并在反向梯度更新的时候会去读写。同理，使用 HashTable 方案时，我们需要同时创建两张单独的 HashTable 用来训练 m、v 参数。那么在前向，反向中需要分别对 Embedding、 m、v 进行一次 Lookup 和一次 Insert，总共需要三次 Lookup 和三次 Insert。

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

图 14 基于 HashTable 的参数融合策略

该特性属于一个普适型优化，开启聚合功能后，训练速度有了显著的提高，性能提升幅度随着模型和 Worker 规模的变化，效果总是正向的。在美团内部真实业务模型上，聚合之后性能相比非聚合方式能提升了 45% 左右。

3.4.2 Embedding 流水线优化

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

• 如果将分布式训练简单的抽象为计算和通信两个过程，绝大多数主流的深度学习框架都支持在执行计算图 DAG 时，通信和计算的 Overlap。

• 如果将深度模型训练简单的分为前向和反向，在单步内，由于两者的强依赖性，无法做到有效并行，字节 BytePS 中引入的通信调度打破了 step iteration 间的屏障，上一轮的部分参数更新完毕后，即可提前开始下轮的前向计算，增强了整体视角下前反向的 Overlap。

• 百度 AIBox 为了解决 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 主要包含稠密部分子网络计算、梯度计算、稠密参数部分更新等环节。

图 15 Embedding 流水线模块交互关系

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

图 16 Embedding 流水线架构流程图

上图是 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 作为底层 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% 左右。

图 17 HashTable 内存优化

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 算子存在冗余数据遍历，这些操作其实可以在一次数据遍历中全部做完，节省掉算子切换、冗余数据遍历的耗时。

图 18 Unique 算子内部出现 DRAM Bound 问题

总结来说，HashTable 开辟过大会导致大量的 minor_page_fault，导致访存的时间增加，HashTable 过小又可能会导致扩容。我们采用了基于启发式算法的内存自适应 Unique 算子实现，通过对训练历史重复率的统计，我们可以得到一个相对合理的 HashTable 大小，来提高访存的性能；另外 Unique 算子内 HashTable 的具体选择上，经过我们的多种测试，选择了 Robin HashTable 替换了原生 TF 中的实现。

进一步，我们对围绕 Embedding ID 的 Unique 和 Partition 环节进行了算子合并，简化了逻辑实现。经过上述的优化，Unique 单算子可以取得 51% 的加速，在真实模型端到端上可以获得 10% 左右的性能提升，算子总数量降低了 4%。

在整个关键算子优化的过程中，Intel 公司的林立凡、张向泽、高明进行大量的技术支持，我们也复用了他们的部分优化工作，在此深表感谢！

大规模稀疏算法建模

大规模稀疏能力在业务落地的过程中，算法层面还需要从特征和模型结构上进行对应升级，才能拿到非常好的效果。其中外卖广告从业务特点出发，引入大规模稀疏特征完成外卖场景下特征体系的升级，提供了更高维的特征空间和参数空间，增强了模型的拟合能力。重新设计了面向高维稀疏场景的特征编码方案，解决了特征编码过程中的特征冲突问题，同时编码过程去掉了部分冗余的特征哈希操作，一定程度上简化了特征处理逻辑，并降低了特征计算的耗时。

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

总结与展望

TensorFlow 在大规模推荐系统中被广泛使用，但由于缺乏大规模稀疏的大规模分布式训练能力，阻碍了业务的发展。美团基于 TensorFlow 原生架构，支持了大规模稀疏能力，并从多个角度进行了深度优化，做到千亿参数、千亿样本高效的分布式训练，并在美团内部进行了大规模的使用。对于这类关键能力的缺失，TensorFlow 社区也引起了共鸣，社区官方在 2020 年创建了 SIG Recommenders，通过社区共建的方式来解决此类问题，美团后续也会积极的参与到社区的贡献当中去。

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

作者简介

逸帆、家恒、峥少、鹏鹏、永宇、正阳、黄军等，来自美团基础研发平台，机器学习平台训练引擎组，主要负责美团分布式机器学习训练系统的性能优化与能力建设。

海涛，来自美团外卖广告策略团队，主要负责美团外卖广告业务的算法探索和策略落地工作。

参考文献

[1] https://www.usenix.org/system/files/conference/osdi16/osdi16-abadi.pdf

原文：TensorFlow在推荐系统中的分布式训练优化实践
中文：TensorFlow 公众号