美团外部深度定制的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申请受权。