TensorFlow与mxnet分布式框架对比

Mxnet和TensorFlow是目前国内最流行的两个开源分布式深度学习训练系统, 两者在设计上有些不同.

mxnet是Parameter Server架构, server和worker是两个最主要的进程,另外还有个负责集群管理的scheduler进程。server负责分布式存储模型参数,worker负责计算,且worker之间不能直接通信,只能通过server互相影响,一般来说,mxnet常用来做数据并行,每个GPU设备上都训练完整的DL模型。TensorFlow主要由client,server,worker三种进程,它虽然也可以实现parameter server的功能,但它计算节点的通信方式并不是只能靠server来完成,TensorFlow的woker是可以互相通信的,可以根据op的依赖关系主动收发数据。

另一个不同点是关于op和device的对应,一般情况下,mxnet常被用来做数据并行,每个GPU设备上都包含了计算图中所有的op。而TensorFlow是可以由用户指定op的放置的,大部分情况下,一个GPU设备只负责某个或者某几个op的训练任务,因此,也催生了关于op placement算法的研究,Google最近发表了一篇使用LSTM模型进行op placement优化的论文,我也在知乎上回答过该论文的详细内容:https://www.zhihu.com/question/61210638/answer/188682910,那个抱着双手的老虎就是我。

这篇笔记主要对比mxnet和TensorFlow在分布式框架下的一些不同,主要包括:进程类型、如何在分布式环境下启集群、如何做初始化。

1.进程类型

mxnet tensorflow
worker client
server server
scheduler worker

2.网络通信库

Mxnet tensorflow
zmq grpc

3.进程间通信模式

Mxnet tensorflow
Worker与scheduler,控制信息 Client与server,控制信息
Server与scheduler,控制信息 Server与worker,控制信息
Worker与server,数据 Worker与worker,数据

4.通信模式

Mxnet tensorflow
消息(队列),PostOffice负责多机之间消息收发,每个机器节点上都有唯一的postoffice单例对象。 消息(std::map),Rendezvous负责多机之间消息收发。每个机器节点上都有唯一的Rendezvous对象。

5,分布式启动

Mxnet tensorflow
各节点执行相同脚本。

根据集群host配置在每个节点上分别启动worker,server,scheduler进程,并进行集群初始化例如worker和server向scheduler的注册等等。

此时集群各个节点之间通信已ready,必要的节点之间可以互相收发数据。

各节点执行相同脚本。

根据集群host配置在每个节点上分别启动client,server,worker进程。完成集群初始化例如GRPC的service启动等。

此时集群各个节点通信已ready,必要的节点之间可以互相收发数据。

  1. 初始化流程6.1 Mxnet:系统初始化是由ps-lite这个模块完成的。每个节点都有个postoffice对象,该对象在系统初始化时进行一系列操作:
    1.     在构造函数中,根据环境变量设置当前节点的
    2.     根据配置文件给worker,server,scheduler进程分配编号。
    3.     调用Van对象初始化各个节点之间的通信,包含以下几个步骤:
      • 获取role为scheduler的ip和端口地址,并建立本节点与scheduler的通信。
      • 启动一个receiveing线程用来接收remote节点的消息。
      • 发送一个message给scheduler报告自己的情况。
      • 启动一个线程用来向scheduler发送heartbeat。

    6.2 TensorFlow:

    1. 启动GRPC Server
    2. 在各个节点上分别启动一个WorkerService线程和一个MasterService线程。
  2. 消息处理

7.1 Mxnet

worker和server以及worker->scheduler、server->scheduler之间通过发送message进行通信。这些消息都会被postoffice这个单例对象负责处理。

在6.3.2中,zmq收到的消息会被放到customer的ThreadsafeQueue<Message>队列中,在customer的Receiving函数中,从消息队列中取出消息并交给recv_handle_处理。而recv_handle_函数就是KVWorker和KVServer的Process函数。

KVWorker::Process函数是worker进程的函数,只是在收到消息后执行回调,例如对pull到的数据进行计算。

KVServer::Process函数是server进程的函数,收到消息之后,调用request_handle_进行处理。request_handle_是一个函数指针,其指向的函数在用户程序中,由用户负责实现。

7.2 Tensorflow

Tensorflow中数据的发送和接收是由send Op和recv Op分别实现的, 发送的消息由Rendezvous类负责处理。

对于send Op, 首先给要发送的消息构造一个unique key,并且交给Rendezvous处理Rendezvous会根据key在自己的消息map中查找是否有对应的recv请求,如果有,就执行recv注册的回调函数;如果没有,就把<key,message>对存放到自己的消息map中。

对于recv Op, 首先给要接收的消息构造一个unique key,并交由Rendezvous处理,Rendezvous根据key在自己的消息map中查找是否已经存在所需的消息,如果已存在,就直接copy消息;如果不存在,就调用GRPC的接口向remote节点发起异步远程调用。

以上所有,都只是关于分布式框架部分的流程,不涉及深度学习模型分布式训练的业务需求,这些以后会记录下来。

2 thoughts on “TensorFlow与mxnet分布式框架对比”

Leave a Reply

Your email address will not be published. Required fields are marked *