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

文章插图

图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，并且对于生产者先于消费者的情况，可以加快消费端数据获取的速度。