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

Posted 美团技术团队

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了TensorFlow在美团外卖推荐场景的GPU训练优化实践相关的知识,希望对你有一定的参考价值。


总第497

2022年 第014篇

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

  • 2 GPU训练优化挑战

  • 3 系统设计与实现

  • 3.1 参数规模的合理化

  • 3.2 系统架构

  • 3.3 关键实现

  • 4 系统性能优化

  • 4.1 数据层

  • 4.2 计算层

  • 4.3 通信层

  • 4.4 性能指标

  • 5 业务落地

  • 5.1 完备性

  • 5.2 训练效果

  • 6 总结与展望

  • 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计算用到的各种数据提前放置到显存中。而从2016年~2020年,NVIDIA Tesla GPU卡[5]计算能力提升了10倍以上,但显存大小只提升了3倍左右。
  • 其它维度资源并不是很充足:相比GPU算力的提升速度,单机的CPU、网络带宽的增长速度较慢,如果遇到这两类资源负载较重的模型,将无法充分发挥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架构的架构图如下所示:
    图1 系统架构
    整个系统主要包括三个核心模块:数据模块,计算模块,通信模块:
  • 数据模块:美团自研了一套支持多数据源、多框架的数据分发系统,在GPU系统上,我们改造数据模块支持了多网卡数据下载,以及考虑到NUMA Awareness的特性,在每颗CPU上都部署了一个数据分发服务。
  • 计算模块:每张GPU卡启动一个TensorFlow训练进程执行训练。
  • 通信模块:我们使用了Horovod[7]来做分布式训练的卡间通信,我们在每个节点上启动一个Horovod进程来执行对应的通信任务。
  • 上述的设计,符合TensorFlow和Horovod原生的设计范式。几个核心模块可以相互解耦,独立迭代,而且如果合并开源社区的最新特性,也不会对系统造成架构性的冲击。
    我们再来看一下整个系统的简要执行流程,每张GPU卡上启动的TensorFlow进程内部的执行逻辑如下图:
    图2 进程内部执行逻辑
    整个训练流程涉及参数存储、优化器、卡间通信等几个关键模块。对于样本的输入特征,我们分为稀疏特征(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到其他卡中,每张卡拿到自己的稀疏梯度后再执行稀疏优化器,完成大规模稀疏特征的优化。详细流程如下图所示:
  • 图3 稀疏特征处理流程
  • 稀疏特征规模较小,使用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的压力会变得更大。我们先看一下在当前系统下的样本数据流程,如下图所示:
    图4 样本数据流程及核心优化点
    核心流程:数据分发进程通过网络读取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结果如下图:
    图5 样本解析profiling结果
    究其原因,CTR场景的训练样本通常包含了大量的int64类型的特征,int64在PB中是以Varint64类型数据存储的,ReadVarint64Fallback方法就是用来解析int64类型的特征。普通的int64数据类型需要占用8个字节,而Varint64针对不同的数据范围,使用了变长的存储长度。PB在解析Varint类型数据时,首先要确定当前数据的长度,Varint用7bit存储数据,高位1bit存储标记位,该标记位表示下一个字节是否有效,如果当前字节最高位为0,则说明当前Varint数据在该字节处结束。我们实际业务场景的ID特征大多是经过Hash后的值,用Varint64类型表达会比较长,这也就导致在特征解析过程中要多次判断数据是否结束,以及多次位移和拼接来生成最终数据,这使得CPU在解析过程中存在大量的分支预测和临时变量,非常影响性能。以下是4字节Varint的解析流程图:
    图6 ProtoBuf 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字节数据):
    图7 ProtoBuf Varint解析优化后流程图
    对于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上之后才能执行模型训练,如下图所示:
    图8 nsys profiling结果
    对于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主要包含稠密部分子网络计算、梯度计算、稠密参数部分更新等环节。
    图9 Embedding流水线模块交互关系
    两张子图的交互关系为: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使用率也非常低,如下图所示:
    图10 nsys profiling结果
    从图中可以看出,训练期间卡间通信耗时比较长,同时在通信期间GPU使用率也非常低,卡间通信是影响训练性能提升的关键瓶颈点。我们对通信过程进行拆解打点后发现,卡间通信(AllToAll、AllReduce、AllGather等)协商的时间远远高于数据传输的时间:
    图11 Horovod timeline结果
    分析具体原因,以负责大规模稀疏参数通信的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的训练性能,具体数据如下:
    图12 CPU/GPU训练吞吐对比
    可以看到,我们优化后的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.

    ----------  END  ----------

    也许你还想看
      | TensorFlow在推荐系统中的分布式训练优化实践
      | 基于TensorFlow Serving的深度学习在线预估
      | 使用TensorFlow训练WDL模型性能问题定位与调优

    阅读更多

    ---
    前端 |  算法 | 后端 | 数据
    安全 | android | iOS  | 运维 | 测试


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


    总第481

    2021年 第051篇

    美团内部深度定制的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 总结与展望

  • 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所示),可以精准定位到性能的瓶颈问题。

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

    图2 自动化实验框架

    2.2.2 业务视角的负载分析

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

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

    图4 扩展PS提升训练性能实验
    导致这种结果的核心原因是:Worker单步训练需要和所有的PS通信同步完成,每增加1个PS要增加N条通信链路,这大大增加了链路延迟(如下图5所示)。而一次训练要执行上百万、上千万步训练。最终导致链路延迟超过了加PS算力并发的收益

    图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所示:

    图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上面,这会带来一个问题:每个优化器只拥有一个和一个,且仅位于某个PS上。因此,在参数优化的过程中,该PS会承受远高于其他PS的请求,从而导致该PS成为性能瓶颈。

    图7 Adam优化算法
    但是通过观察Adam的优化算法,我们可以看到都是常量,且蓝色高亮的部分都是相对独立的计算过程,各个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单边通信模式。

    图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[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主要包含稠密部分子网络计算、梯度计算、稠密参数部分更新等环节。

    图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[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%左右。

    图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公司的林立凡、张向泽、高明进行大量的技术支持,我们也复用了他们的部分优化工作,在此深表感谢!
    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 参考文献
    [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

    ----------  END  ----------
    招聘信息
    美团机器学习平台大量岗位持续招聘中,社招/校招均可(欢迎投递我们的校招北斗岗位:美团机器学习平台基础架构),坐标北京/上海,构建多领域的公司级机器学习平台,帮大家吃得更好,生活更好。简历可投递至:huangjun03@meituan.com
    也许你还想看
      | 一站式机器学习平台建设实践
      | 美团深度学习系统的工程实践
      | 基于TensorFlow Serving的深度学习在线预估

    阅读更多

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



    以上是关于TensorFlow在美团外卖推荐场景的GPU训练优化实践的主要内容,如果未能解决你的问题,请参考以下文章

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

    Apache Kylin在美团数十亿数据OLAP场景下的实践

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

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

    Flutter Web在美团外卖的实践

    大规模异构图召回在美团到店推荐广告的应用