TensorFlow 大规模稀疏模型异步训练的分布式优化

文 / 刘童璇,Alibaba PAI 团队

Alibaba PAI 团队从 16 年开始在 TensorFlow 上进行优化,结合阿里巴巴推荐、搜索、广告等核心业务,打造锤炼 TensorFlow 对超大规模稀疏模型的训练能力。在 2018 年 9 月在上海举行的” 谷歌开发者日”,和 2019 年 3 月在美国举行的 “TensorFlow Developer Summit” 都曾分享过 (PAI: Platform of A.I. in Alibaba)。后续会持续和 TensorFlow 团队保持紧密的合作,通过 RecSys SIG,逐步开源内部针对超大规模稀疏模型训练的核心能力。

背景

原生 TensorFlow 中使用 PS-Worker 模式进行异步训练时,通过开源的 gRPC 通信库进行不同节点之间的数据同步。超大规模稀疏模型训练通常需要几百到上千 Worker 节点,此时开源的 gRPC 通信库性能存在明显的瓶颈。特别是稀疏模型中通常包含上百的特征列,带来了大量的小包通信,而 gRPC 库的小包通信能力比较差。在原生 TensorFlow 分布式框架下,难以进行超大规模的稀疏模型训练。

针对原生 TensorFlow 中进行超大规模稀疏训练时分布式扩展性不足的问题,Alibaba PAI 团队为 TensorFlow 社区贡献了 grpc+seastar 及 FuseRecv 两个功能。在一些典型的业务场景下,大大提升了稀疏模型训练的分布式扩展性。TensorFlow 社区内部分同学试用了这两个功能,在 400 worker 规模下能够提升 2-4 倍。

关键技术

通信算子融合 (FuseRecv)

原生 TensorFlow 中异步训练中,使用的 Send/Recv Op 进行跨节点的 Tensor 传输,其中 Recv 节点对应一个 RPC 请求来传输一个 Tensor,即使有 Recv Ops 输出到相同的目标 Op。而且每个即使接收到的张量是标量,Recv 节点也会触发一个 RPC 操作。如图:

image1

为了降低 RPC 的请求个数,同时确保不会因为 RPC 合并带来计算通信 Op 的无法很好 Overlap 问题。在图优化阶段,基于图拓扑进行 Recv Op 的合并,将多个 RecvOp 合并为 FuseRecvOp。RPC 在 WorkerInterface 增加 FuseRecvTensorAsync 接口,支持一个 RPC 传输多个 Rendezvous Key 和获取多个 Tensor。

如图所示,相比原生 TensorFlow 中 a, x 需要两个 Recv 节点进行两次 RPC,FuseRecv 节点会包含两个 slot,通过一次 RPC 将 a, x 从 Device A 中的 Rendezvous 中获取。

image3

您可参考 FuseRecv 的 详细设计文档,目前 RFC 已经过社区评审,被 TensorFlow 社区接受。我们正在将代码贡献给社区。

用户态零拷贝的数据传输

为了追求更好的性能,我们重新设计了 Wire Format,避免了通信过程中的序列化、拷贝的开销。此外我们打通了 RPC 层的内存管理和 TensorFlow 的内存管理,从而实现了 Tensor 的零拷贝传输。

grpc+seastar 的 Request 及 Response 均由定长的 Header 附带一段或多段 Payload 组成。Header 与 Payload 可以独立的进行内存分配,并按照给定的顺序发送,从而在网络上形成完成的报文字节流。

在 RPC 的设计中,一个重要的问题是要在 wire format 层面解决 Response 消息与 Request 消息的关联。一种比较容易的想到的做法是将每次 RPC 请求用一个自增的 id 进行标识,然后再通过类似 map 的数据结构将 id 与 RPC 请求上下文进行关联。每次 Request 及 Response 的 wire format message 结构中,都需要携带该 RPC id,从而关联至上层逻辑。对于这种设计来说,一个显而易见的问题即是基于 RPC id 寻找上下文的 “查表” 开销会随着并发请求量的增大而迅速上涨,尤其在 TensorFlow 这种异步多线程场景下,很难做到无锁操作。

为了避免上述问题,我们采用了直接指针消息映射的策略:在我们的 Request 的 RPC Header 中,会有字段专门存放 context 的指针,而这一指针会在 Response 的 RPC Header 中被原样回传,从而实现直接内存映射,我们以 RPC client 端的场景为例,展示详细的数据结构如下图所示:

在 Tensor 的接收端,在完成定长的 tensor meta 字段的接收过程后,接收端的代码会直接调用对应 TensorFlow 中 Device 的 Allocator 为待接收的 Tensor 分配出适当的内存空间,底层通信层可以直接将 Tensor 的数据接收到该内存空间内。

通信收发链路无锁化

Seastar 是开源的通信库 (https://github.com/scylladb/seastar) 其具有的 Sharing Nothing 架构、BusyPolling 机制、 Message-Passing By Queue 的设计能够大大降低共享对象的锁开销、context switch。特别是在小包通信的场景下,Seastar 展现了数倍于 gRPC 的通信能力。

我们基于 Seastar 的设计特点,在通信收发链路的设计中延续了 Share-nothing 的设计,在完整的通信收发链路中整个流程通过 tls 实现了多线程的 share-nothing,并且基于 Message-Passing By Queue 的方式和外部的 TensorFlow 工作线程进行交互。

目前代码已经贡献给 TensorFlow 社区,请参考 RFC代码设计文档

效果

下面是两个测试模型的性能数据。从下图来看 gRPC 从 100 个节点规模往上就就很难继续提高计算吞吐,在 W&DL 模型下从 100 个节点开始反而整体的计算能力开始下降。而使用 grpc+seastar 和 FuseRecv 功能后,则能够继续扩展,提供更高的计算能力。

经验分享

这是我们在 TensorFlow 社区推进的较大规模的 RFC,学到的一些经验包括:

  1. 提供良好的覆盖很多细节的设计文档,可以帮助沟通,特别是比较少写英文文档的同学。
  2. 可以找相关 TensorFlow 团队同事,线下提前沟通,在提交给社区前有一些初步意见和迭代。而提交 RFC 给社区的时候,强有力的 TensorFlow 团队 Sponsor,也可以加速此过程。
  3. TensorFlow 有很多 SIG 组织,一些贡献可以考虑通过 SIG 组织来获取反馈,比如 SIG AddOns,SIG IO,SIG Networking 等。grpc+seastar 的贡献过程得到了 SIG Networking 的 Bairen Yi 和 Jeroen Bédorf 的大力支持,给了很多建议。
  4. 获取社区中其他开发者的反馈也很有帮助。在以上功能的开发过程中,社区中其他同学提前尝试了我们的功能,并给了很好的反馈。

致谢

感谢在 grpc+seastar,FuseRecv 功能 merge 进社区过程中,给过很多建议和帮助的同学们:Shuangfeng Li、Wei Wei、Ayush Dubey、Jeroen Bédorf、Derek Murray、Bairen Yi 和 Paul Tucker。

原文:TensorFlow 公众号

w&d 模型 100worker 下加速了 2 倍吗? 单个 worker 多大规模呢
主要实现了 zerocopy, 无锁化 和 FuseRecv,对性能影响最大的是不是 fuserecv 这块呢?