IJPP 2019 | Improving the Performance of Distributed MXNet with RDMA

MXNet是一款非常优秀的深度学习框架,在工业界和学术界都广受好评。但是,MXNet最初是使用proxy-socket接口实现分布式训练,这导致它在分布式环境中表现平平。在分布式训练过程中,参数更新非常频繁,网络逐渐成为影响训练性能的主要因素。本文描述了一种有效的系统设计方案,即通过基于RDMA的参数服务器接口扩展MXNet。通过对内存使用和传输开销的适当优化,基于RDMA的MXNet比原始版本的性能更好。

Background Review

MXNet Architecture

MXNet是一个基于数据流的深度学习系统,它依赖于后端引擎进行资源管理和任务调度。每个资源单元使用唯一标签注册到引擎,并且通过指定的资源标签将所有操作(矩阵操作,数据通信)推入引擎。系统以异步多线程模式工作,以实现数据并行和模型并行。

图1是MXNet通信系统的多层结构。在通信系统的高层,MXNet通过全局的KVStore与参数服务器(PS-Lite)打交道;在中层,PS-Lite通过ZMQ实现数据传输;底层的ZMQ则通过TCP/IP协议进行网络传输。通过基本的RDMA语义模型,本文提出了一种简单的设计方案,将TCP/IP替换为RDMA。

Parameter Server Framework

参数服务器架构中,主要有两种角色,worker和server(有时也会有scheduler)。worker负责本地计算,而server则维护全局参数。每次迭代时,worker首先读取训练数据,然后从server拉取参数,再计算出本地梯度,最终把本地梯度推送到server。与此同时,server将会聚合所有worker推送的梯度,并更新全局参数。在MXNet中,参数服务器通过PS-Lite实现。

PS-Lite针对机器学习应用的特点设计了两种通信机制:request/response和push/pull。前者用于交换scheduler的控制消息,后者用于在节点之间传输数据消息。需要注意的是,网络拥塞主要由频繁的数据通信造成的,所以push/pull操作是通信开销的主要来源。因此,后文主要针对push/pull进行优化。

RDMA and InfiniBand

InfiniBand是一种开放式工业标准交换结构,目的是互连HPC集群中的节点。InfiniBand最大的特点就是RDMA,它可以让一台机器直接读写另一台机器的内存而不经过CPU。InfiniBand提供两种信息传递语义:内存语义和消息语义。内存语义引入RDMA写和RDMA读操作;而消息语义中,发送/接收操作被用于通信。RDMA是一种单边操作,不会在远程产生开销。使用RDMA必须要事先知道目标地址,且必须在接收方对传入的消息进行显式检测。

IB谓词是一种在InfiniBand上支持RMDA的底层通信接口。用户通过这些接口来创建、管理和销毁资源。IB谓词包括QP(queue-pairs)、CQ(completion queues)、MR(memory regions)和PD(protection domains)。

Contributions

Communication Framework

MXNet的通信框架完美地结合了KVStore、PS-Lite和ZMQ,并且在执行模型、程序逻辑和高性能之间取得平衡。MXNet的详细通信架构如图2所示。

首先,工作负载被分配到工作节点中的多台计算设备上,因此我们会从不同的计算设备上接受到许多中间结果。举例来说,每个设备执行多次迭代,并行地生成独立的键值对。KVStore并不推送这些原始地键值对,而是通过相同的键将中间值合并到本地聚合结果中,如图2第一步所示。紧接着,PS-Lite将聚合后地键值对分成多个段,分段策略和目标由服务器的编号和全局键地范围预定义。最终,ZMQ负责完成数据传输,将分好的段传输给全局共享的KVStore。

我们可以这样理解KVStore、PS-Lite和ZMQ的关系:KVStore相当于分布式数据库,用于存储MXNet的参数;PS-Lite负责分布式深度学习中的通信逻辑;而最底层的ZMQ则为数据的高速传输提供支持。

值得注意的是,整个通信系统对计算系统来说是透明的。也就是说,无论使用CPU还是GPU计算,本文所提出的优化方法都是适用的。

Basic Model

前面提到,MXNet的通信系统由KVStore、PS-Lite和ZMQ三级组成。

Memory Reuse

Batch-based Transfer