TensorFlow在美团推荐系统中的分布式训练优化实践( 四 ) _云知道

该优化所带来的提升具备普适性且效果明显，在美团内部某业务模型上，通过β热点去除可以带来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常有几百上千个，这会带来新的问题：