TensorFlow在美团推荐系统中的分布式训练优化实践( 三 ) _知识分享

但是通过观察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静态分配器

这里核心的设计思路为：
虽然稀疏场景同一个算子输出Tensor的Shape存在变化的可能，但是整体变化幅度可控，通过监控与分析，是可以找到一个较为稳定的内存大小，满足多Step间Tensor的存储需求。基于上面的信息，我们修改了原有逐Tensor(Request)的MR申请策略，通过一次性预申请一块较大的空间并注册到网卡端，后续通过自己维护的分配策略进行空间的分配，大大降低了MR申请的频率，绝大多数情况下，训练全过程中只需要一次MR注册申请即可。我们引入了一种简单的交换协议，将传输Tensor的Shape，Data打包到一起，写到Client端。Client端根据协议，解析出Tensor大小，并最终读取Data，避免了原生实现中因Tensor的Shape变化而产生的多次协商过程。