揭秘滴滴机器学习平台参数服务器 -- 来自工业界的性能基准

在机器学习领域，分布式训练是一个持久的课题。单节点硬件性能的提升通常难以满足业务不断提高的算力需求，而分布式训练一个重要目标就是做到水平扩展，提供更高的整体算力，在尽量短的时间内处理更多的数据，或者进行更多的尝试以获得精度更高的模型。在滴滴的机器学习平台，无论是在实验环境进行模型调参迭代，还是在生产环境进行日常的线上模型训练更新，分布式训练不仅保证了时效性，提高了效率，而且为自动调参，自动求解这类非常规手段提供了基础支持。

目前，很多机器学习框架都给出了分布式训练的解决方案。相对于业界其他厂商，滴滴分布式训练所面临的挑战显然更大一些。从数据上来讲，滴滴机器学习平台上所服务的业务不但涵盖了传统的图像、语音、视频等业务，并且面临着来自于交通、地图、路况等领域更大规模的挑战。以北京为例，进行 ETA 模型训练(到达时间预估)一天的训练数据的大小就超过了所有公开图像数据集；而总体上每天滴滴订单量超过 3000 万单，每天产生交通数据量极为庞大。因此对滴滴来说，分布式训练面临的挑战巨大。

在分布式训练中，网络传输数据量大，因此参数同步的性能就显得非常重要。而在分布式训练中，参数服务器(Parameter Server)扮演着参数同步的角色，其同步性能直接影响着分布式训练的整体性能。在目前的实践中，学术界，工业界往往将关注点放在了梯度压缩、异步梯度下降等方法来提高性能，而对底层的参数同步机制和性能关注不够，只被动的使用现成框架或业界已有的解决方案。而如果能够从底层提高参数服务器的同步性能，则能够普惠上层各种算法、技巧。

目前业界参数同步方案一般采取 Allreduce 来完成；较为普遍采用的分布式通信框架有 OpenMPI、NVIDIA 公司的 NCCL2 等。MPI(Message Passing Interface)是工业界的一个标准，一个跨语言的通信协议，用于并行计算等方面；而 OpenMPI(Open MPI)作为高性能消息传递库，是 MPI 标准的开源实现。其实现了点对点以及 collective 的通信语义，支持 TCP、RDMA 等网络传输方式，侧重于高性能。而 NCCL(NVIDIA Collective Communications Library)是 NVIDIA 公司为多 GPU 开发的用于 GPU 之间实现 collective 通信的库。号称高效、高性能的实现了 GPU 之间信息交换和通信。很多参数服务器实现在底层利用了 MPI 或 NCCL 的 Allreduce API，比如 Uber 的 Horovod。我们实际测试发现：无论是 OpenMPI 还是 NVIDIA 的 NCCL2 的 Allreduce API，其性能虽然不错，但是离理想水平相差甚远，不尽人意，都有很大的优化空间。

在滴滴内部，目前使用最多的深度学习框架是谷歌的 Tensorflow。它虽然提供了分布式训练的解决方案，但方案里使用 gRPC 来通信，基于 TCP 协议，性能不高，且可用性也不好，API 使用较为复杂。Tensorflow 也提供了 RDMA 的传输方式，但 Tensorflow 设计的时候并没有考虑底层传输对 RDMA 的支持，因此其设计方案并不能很好的贴切地发挥 RDMA 传输的性能优势；反而是其 RDMA 实现更多的去适配和套入 Tensorflow 本身低效的 gRPC 通信框架。这类设计往往是代码抽象上考虑的很多，但是实现性能上考虑的很少，削足适履。Tensorflow RDMA 方案不合理的地方举例如下：

1.Tensorflow 分配要传输或接收的 tensor 内存时候采取了动态分配，这就非常不利于 RDMA 传输，RDMA 传输进行之前需要注册内存，而注册内存的调用性能代价很大，动态分配内存意味着频繁的注册内存，性能开销很大。因此，Tensorflow RDMA 实现采取了 HKUST 0 copies 方案来缓解这个问题；

2.为了适应 tensorflow gRPC 通信的 Rendezvous 协议，RDMA 建立一个中间层 cache 来存放 tensor 的一些元信息，比如 shape, data-type 等，并放在了 sender 端，而 Rendezvous 协议要求 receiver 主动发起传输，则 receiver 首选需要拿到这些元信息才能通信。并且，过程中元信息还有可能更新。这些都是额外的通信开销。

3.Tensorflow RDMA 通信的通道不可以一次发送多个消息，本次 RDMA_WRITE 完成前，通道不可以再被使用；因此 Tensorflow 采用了两个状态变量来控制和约束，增加了额外的处理和负担。

类似上述的细节地方，还有很多；有兴趣的读者可以参阅 Tensorflow RDMA 实现的官方文档或代码；另外值得一提的是，哪怕最新版 1.8 版本的 Tensorflow，其官方的二进制安装源依然默认不支持 RDMA，需要自己从源码手动编译来支持 ，这也是谷歌对自己的 RDMA 实现并没有信心的表现。

在开源的参数服务器方案中，Uber 的 Horovod 有许多人使用，官方号称的性能也很高。但它底层其实是采用 MPI 或 NVIDIA 的 NCCL 来完成最底层的 Allreduce 同步梯度。如前面所述，MPI 和 NCCL 的 Allreduce 性能，在我们滴滴看来并不算是最高效的实现；我们有能力来构建出更加高效的参数服务器实现来完成参数同步。因此，在了解到业界方案的不足之后，滴滴机器学习平台研发了基于 RDMA 网络的参数服务器解决方案 DiDi PS，突出性能和效率。

DiDi PS 采取环状结构，基于 ibverbs API 实现了高效的 RDMA 通信的 Allreduce 算法。当然，在优化中我们也发现厂家的 libverbs 从性能角度也有不少可以优化的地方，这个以后可以再开专题讨论。 环状结构而非中心集中的 server-client 模式，减少了通信量，同时可以消除网络传输可能的带宽竞争和网络拥塞。底层自研的高效 RDMA 通信库，以独特的方式规避了 ibverbs 的 RDMA 通信开销，充分挖掘了硬件的属性来提高性能。另外，我们自研的 Allreduce 算法细致地重叠了计算和传输，尽量减少了不必要的内存拷贝来减少额外代价；最后我们也充分考虑了 GPU 拓扑、CPU 亲和性等硬件属性来提高性能。

在我们机房的实际环境下，我们对比测试了 DiDi PS 和业界其他的解决方案性能。我们首先对比的是 Allreduce 这一基本原语的性能，对标 OpenMPI 和 NVIDIA 的 NCCL2。

测试环境介绍如下：

我们采用 256M Bytes 的单精度求平均所用的时间来对比 Allreduce 的性能，分别测试 5 机器 20 节点、4 机器 16 节点、3 机器 12 节点、2 机器 8 节点等不同规模下 allreduce 性能。测试中保证每个节点（每个进程）可以使用一张 GPU 卡，没有 GPU 资源的竞争。实际测试的结果如下图所示：

通过上述测试可以发现，我们 DiDi PS 的 Allreduce 性能，相对于 OpenMPI 和 NVIDIA 的 NCCL2，均处于绝对领先优势。

为了便于内部团队在 Tensorflow 内使用 DiDi PS，我们将 DiDi PS 的基本操作进行了封装，成为了 Tensorflow 的 Operation，可以像 Tensorflow 原生的 Operation 一样方便使用。对此，为了进一步充分测试，我们选择使用实际的训练模型来对比 DiDi PS 和 Tensorflow distributed training 方案，以及 Uber 的 Horovod 的训练性能。这里训练模型采用 inception v3 模型，在实际的 imagenet-data 上进行训练。Tensorflow 采用官方的 1.8.0 版本，RDMA 的 Tensorflow 采取自己基于 1.8.0 版本的代码编译二进制完成安装；Uber 的 Horovod 采用最新版本 0.13.2。实际测试结果如下图所示：

由上图可见，DiDi PS 在实际模型的训练对比中，性能也领先于 Tensorflow 以及 Uber 的解决方案。另外在测试中我们发现，Horovod 底层采用 NCCL2，采用 RDMA 通信的时候，性能很不稳定，波动很大。同样，Tensorflow 底层 RDMA 方案测试过程中也不稳定，波动甚至程序奔溃退出过几次。可见，在稳定性方面，DiDi PS 也具有领先优势。

另，DiDi PS 虽然目前做到了不错的性能，但我们依然还在持续的优化当中，期待后续在取得更大的进步，为业务需求提供基础支持。

本文转载自公众号滴滴技术（ID：didi_tech）。

原文链接：

https://mp.weixin.qq.com/s/LiiKBAw_d5lLRSoTe6pYcQ