[RFC] Unified API for Distributed Data Parallel Training

[eric-haibin-lin]

Background

Data parallel training is the most common distributed training technique when it comes to multiple GPUs or multiple hosts. Currently, several communication backends provide functionalities for communicating tensors across devices/hosts for data parallel training. For MXNet users, there are a few options:

1. native kvstore
2. p3 kvstore
3. horovod
4. bytePS

These different implementations provide different APIs:

• native kvstore
  • high level APIs: mx.gluon.Trainer
  • low level APIs: kv.push, kv.pull, kv.init
• horovod
  • high level APIs: hvd.init(), hvd.DistributedTrainer
  • low level APIs: hvd.broadcast, hvd.allreduce
• bytePS
  • high level APIs: bps.init(), bps.DistributedTrainer
  • low level APIs: byteps_declare_tensor, byteps_push_pull

Here, high level APIs refers to the API a typical novice user uses for a distributed training job. To communicate tensors not managed by Trainer or DistributedTrainers, users may refer to the low level APIs to send/receive a custom tensor.

Problem Statement

Sometimes we want to easily switch between these different distributed communication backends and compare which one performs the best for a particular distributed training environment. Due to different APIs of these implementations, it requires lots of user code changes to try each one of them. It typically involves custom logics to:

1. launch python processes for distributed training job (BytePS launch v.s. horovod)
2. initialize communication backends (example code)
3. create (Distributed)Trainers (example code)
4. send custom tensors (example code)

Proposal

My proposal is to provide a unified API to allow custom communication backends as plugins for MXNet, so that no new user code is required to switch between these backends.

Specifically, communication backend provider implements the following python APIs.

class AbstractKVStore:

• def init(): initialization
• def broadcast(name, tensor, root_rank): broadcast the tensor at root_rank to all ranks
  • name: tensor name. int or str
  • tensor: ndarray
• def push_pull(name, tensor, output): push tensor and pull in output. When optimizer is not set, it performs summation of tensor from all ranks. The result of the summation is then pulled back to output tensor.
  • name: tensor name. int or str
  • tensor: ndarray to push
  • output: ndarray to store the output of pull
• def set_optimizer(optimizer): set the optimizer at parameter servers. Optional interface, only used for parameter server based backends.
  • optimizer: mx.optimizer.Optimizer

A communication backend provider can implement these APIs and register a new KVStore in MXNet via mx.kv.register(). For MXNet users, they only need to interact with the following MXNet APIs:

• using high level APIs for a typical data parallel model training

backend = mx.kv.create('horovod')
trainer = mx.gluon.Trainer(kv=backend)
# forward: loss = net(data)
# backward: loss.backward()
# update: trainer.step()

• using low level APIs to reduce a custom tensor

kv.broadcast("name", custom_ndarray, root_rank=0)
kv.push_pull("name", custom_ndarray, out=custom_ndarray)

Limitation

The unified interfaces do not advanced features such as sparse ndarrays or gradient compression, which is less mature and not provided by all communication backends.

The above proposal targets use case 2,3,4 in the problem statement. It can be extended to tackle 1 as well if the feedbacks are positive.

@ymjiang @apeforest @anandj91 @rich-junwang

[apeforest]

In the Limitation, I suppose you meant 'use case 1,3,4', right?

[eric-haibin-lin]

I did mean use case 2,3,4.
Initialization is done in the constructor kv.__init__(), and for horovod it could be simply a hvd.init() call.

I have not discussed problem 1 for too much details. horovod uses mpirun to setup connection and launch processes, while byteps/p3 and native kvstore currently use the dmlc/launcher script. I do see that dmlc/launcher has mpi support, but I need to play more with it to see if it fits existing use cases. But I don't see fundamental blockers for (1).

[leezu]

Would it make sense to add optional support for sparse ndarrays and gradient compression in AbstractKVStore? You mentioned not all frameworks support it. Do you expect the API to change in the future?

[eric-haibin-lin]

I do expect the API to change in the future. Currently @szhengac @zhongyuchen and I are exploring APIs for gradient compression with a few algorithms, and we may bring back the best practice back to MXNet.