美团机器学习平台基于外部深度定制的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类型为例，咱们的优化流程次要包含两步：
- SIMD寻找最高位：通过SIMD指令将Varint类型数据的每个字节与0xF0做与运算，找到第一个后果等于0的字节，这个字节就是以后Varint数据的完结地位。
- 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申请受权。