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

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变化而产生的多次协商过程。

图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在美团推荐系统中的分布式训练优化实践】TensorFlow社区版的RDMA通信过程，不仅仅包含上面Tensor数据的发送和接收过程，还包括传输相关的控制消息的发送和接收过程，控制消息的发送和接收过程同样是使用了ibv_post_send和ibv_post_recv原语。原生的控制流实现存在一些瓶颈，在大规模训练时会限制控制流的吞吐，进而影响数据收发的效率。具体体现在：
请求的发送通过同一片RequestBuffer内存进行写出，多个Client的请求均依赖这一片Buffer ，也就导致到控制流信息实际是串行发送的，只有等到对端的Ack信息后，才可以下一个Request的写出，限制了请求的发送吞吐。在Client端需要轮询RDMA Completion Queue来获得请求的到达，以及相关状态的变更。原生实现仅有一个Completion Queue ，单线程进行轮询处理，在大规模分布式训练中，限制了应答的效率。