Types of Processes¶

MXNet has three types of processes which communicate with each other to accomplish training of a model.

• Worker: A worker node actually performs training on a batch of training samples. Before processing each batch, the workers pull weights from servers. The workers also send gradients to the servers after each batch. Depending on the workload for training a model, it might not be a good idea to run multiple worker processes on the same machine.
• Server: There can be multiple servers which store the model’s parameters, and communicate with workers. A server may or may not be co-located with the worker processes.
• Scheduler: There is only one scheduler. The role of the scheduler is to set up the cluster. This includes waiting for messages that each node has come up and which port the node is listening on. The scheduler then lets all processes know about every other node in the cluster, so that they can communicate with each other.

KV Store¶

MXNet provides a key-value store, which is a critical component used for multi-device training. The communication of parameters across devices on a single machine, as well as across multiple machines, is relayed through one or more servers with a key-value store for the parameters. Each value in this store is represented by a key and value, where each parameter array in the network is assigned a key, and value refers to the weights of that parameter array. Workers push gradients after processing a batch, and pull updated weights before processing a new batch. We can also pass in optimizers for the KVStore to use while updating each weight. Optimizers like Stochastic Gradient Descent define an update rule, essentially a mathematical formula to compute the new weight based on the old weight, gradient, and some parameters.

If you are using a Gluon Trainer object or the Module API, it uses a kvstore object internally to aggregate gradients from multiple devices on the same machine as well as across different machines.

Although the API remains the same whether or not multiple machines are being used, the notion of kvstore server exists only during distributed training. In this case, each push and pull involves communication with the kvstore servers. When there are multiple devices on a single machine, gradients from these devices are first aggregated on the machine and then sent to the servers.Note that we need to compile MXNet with the build flag USE_DIST_KVSTORE=1 to use distributed training.

The distributed mode of KVStore is enabled by calling mxnet.kvstore.create function with a string argument which contains the word dist as follows:

kv = mxnet.kvstore.create(‘dist_sync’)

Refer KVStore API for more information about KVStore.

Distribution of Keys¶

Each server doesn’t necessarily store all the keys or parameter arrays. Parameters are distributed across different servers. The decision of which server stores a particular key is made at random. This distribution of keys across different servers is handled transparently by the KVStore. It ensures that when a key is pulled, that request is sent to the server which has the corresponding value. If the value of some key is very large, it may be sharded across different servers. This means that different servers hold different parts of the value. Again, this is handled transparently so that the worker does not have to do anything different. The threshold for this sharding can be controlled with the environment variable MXNET_KVSTORE_BIGARRAY_BOUND. See environment variables for more details.

Split training data¶

When running distributed training in data parallel mode, we want each machine to be working on different parts of the dataset.

For data parallel training on a single worker, we can use mxnet.gluon.utils.split_and_load to split a batch of samples provided by the data iterator, and then load each part of the batch on the device which will process it.

In the case of distributed training though, we would need to divide the dataset into n parts at the beginning, so that each worker gets a different part. Each worker can then use split_and_load to again divide that part of the dataset across different devices on a single machine.

Typically, this split of data for each worker happens through the data iterator, on passing the number of parts and the index of parts to iterate over. Some iterators in MXNet that support this feature are mxnet.io.MNISTIterator and mxnet.io.ImageRecordIter. If you are using a different iterator, you can look at how the above iterators implement this. We can use the kvstore object to get the number of workers (kv.num_workers) and rank of the current worker (kv.rank). These can be passed as arguments to the iterator. You can look at example/gluon/image_classification.py to see an example usage.

Updating weights¶

KVStore server supports two modes, one which aggregates the gradients and updates the weights using those gradients, and second where the server only aggregates gradients. In the latter case, when a worker process pulls from kvstore, it gets the aggregated gradients. The worker then uses these gradients and applies the weights locally.

When using Gluon there is an option to choose between these modes by passing update_on_kvstore variable when you create the Trainer object like this:

trainer = gluon.Trainer(net.collect_params(), optimizer='sgd',
                        optimizer_params={'learning_rate': opt.lr,
                                          'wd': opt.wd,
                                          'momentum': opt.momentum,
                                          'multi_precision': True},
                        kvstore=kv,
                        update_on_kvstore=True)

When using the symbolic interface, it performs the weight updates on the server without the user having to do anything special.

Different Modes of Distributed Training¶

Distributed training itself is enabled when kvstore creation string contains the word dist.

Different modes of distributed training can be enabled by using different types of kvstore.

• dist_sync: In synchronous distributed training, all workers use the same synchronized set of model parameters at the start of every batch. This means that after each batch, the server waits to receive gradients from each worker before it updates the model parameters. This synchronization comes at a cost because the worker pulling parameters would have to wait till the server finishes this process. In this mode, if a worker crashes, then it halts the progress of all workers.
• dist_async: In asynchronous distributed training, the server receives gradients from one worker and immediately updates its store, which it uses to respond to any future pulls. This means that a worker who finishes processing a batch can pull the current parameters from server and start the next batch, even if other workers haven’t finished processing the earlier batch. This is faster than dist_sync because there is no cost of synchronization, but can take more epochs to converge. The update of weights is atomic, meaning no two updates happen on the same weight at the same time. However, the order of updates is not guaranteed. In async mode, it is required to pass an optimizer because in the absence of an optimizer kvstore would replace the stored weights with received weights and this doesn’t make sense for training in asynchronous mode. Hence, when using Gluon with async mode we need to set update_on_kvstore to True.
• dist_sync_device: Same as dist_sync except that when there are multiple GPUs being used on each node, this mode aggregates gradients and updates weights on GPU while dist_sync does so on CPU memory. This is faster than dist_sync because it reduces expensive communication between GPU and CPU, but it increases memory usage on GPU.
• dist_async_device : The analogue of dist_sync_device but in asynchronous mode.

Gradient Compression¶

When communication is expensive, and the ratio of computation time to communication time is low, communication can become a bottleneck. In such cases, gradient compression can be used to reduce the cost of communication, thereby speeding up training. Refer Gradient compression for more details.

Note: For small models when the cost of computation is much lower than cost of communication, distributed training might actually be slower than training on a single machine because of the overhead of communication and synchronization.

Setting up the Cluster¶

An easy way to set up a cluster of EC2 instances for distributed deep learning is by using the AWS CloudFormation template. If you can not use the above, this section will help you manually set up a cluster of instances to enable you to use ssh for launching a distributed training job. Let us denote one machine as the master of the cluster through which we will launch and monitor the distributed training on all machines.

If the machines in your cluster are a part of a cloud computing platform like AWS EC2, then your instances should be using key-based authentication already. Ensure that you create all instances using the same key, say mxnet-key and in the same security group. Next, we need to ensure that master has access to all other machines in the cluster through ssh by adding this key to ssh-agent and forwarding it to master when we log in. This will make mxnet-key the default key on master.

ssh-add .ssh/mxnet-key
ssh -A user@MASTER_IP_ADDRESS

If your machines use passwords for authentication, see here for instructions on setting up password-less authentication between machines.

It is easier if all these machines have a shared file system so that they can access the training script. One way is to use Amazon Elastic File System to create your network file system. The options in the following command are the recommended options when mounting an AWS Elastic File System.

sudo mkdir efs && sudo mount -t nfs4 -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2 NETWORK_FILE_SYSTEM_IP:/ efs

Tip: You might find it helpful to store large datasets on S3 for easy access from all machines in the cluster. Refer Using data from S3 for training for more information.

Using Launch.py¶

MXNet provides a script tools/launch.py to make it easy to launch distributed training on a cluster with ssh, mpi, sge or yarn. You can fetch this script by cloning the mxnet repository.

git clone --recursive https://github.com/apache/incubator-mxnet

cd example/gluon/

python image_classification.py --dataset cifar10 --model vgg11 --epochs 1

../../tools/launch.py -n 3 -H hosts --launcher ssh python image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync

../../tools/launch.py -n 3 -H hosts --launcher ssh --sync-dst-dir /tmp/mxnet_job/ python image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync

while read -u 10 host; do ssh -o "StrictHostKeyChecking no" $host "pkill -f python" ; done 10

Manually Launching Jobs¶

If for some reason, you do not want to use the script above to start distributed training, then this section will be helpful. MXNet uses environment variables to assign roles to different processes and to let different processes find the scheduler. The environment variables are required to be set correctly as follows for the training to start:

• DMLC_ROLE: Specifies the role of the process. This can be server, worker or scheduler. Note that there should only be one scheduler. When DMLC_ROLE is set to server or scheduler, these processes start when mxnet is imported.
• DMLC_PS_ROOT_URI: Specifies the IP of the scheduler
• DMLC_PS_ROOT_PORT: Specifies the port that the scheduler listens to
• DMLC_NUM_SERVER: Specifies how many server nodes are in the cluster
• DMLC_NUM_WORKER: Specifies how many worker nodes are in the cluster

Below is an example to start all jobs locally on Linux or Mac. Note that starting all jobs on the same machine is not a good idea. This is only to make the usage clear.

export COMMAND='python example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=scheduler DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND

For an in-depth discussion of how the scheduler sets up the cluster, you can go here.

For tuning performance¶

• MXNET_KVSTORE_REDUCTION_NTHREADS Value type: Integer Default value: 4 The number of CPU threads used for summing up big arrays on a single machine This will also be used for dist_sync kvstore to sum up arrays from different contexts on a single machine. This does not affect summing up of arrays from different machines on servers. Summing up of arrays for dist_sync_device kvstore is also unaffected as that happens on GPUs.
• MXNET_KVSTORE_BIGARRAY_BOUND Value type: Integer Default value: 1000000 The minimum size of a big array. When the array size is bigger than this threshold, MXNET_KVSTORE_REDUCTION_NTHREADS threads are used for reduction. This parameter is also used as a load balancer in kvstore. It controls when to partition a single weight to all the servers. If the size of a single weight matrix is less than this bound, then it is sent to a single randomly picked server; otherwise, it is partitioned to all the servers.
• MXNET_ENABLE_GPU_P2P GPU Peer-to-Peer communication Value type: 0(false) or 1(true) Default value: 1 If true, MXNet tries to use GPU peer-to-peer communication, if available on your device. This is used only when kvstore has the type device in it.

Communication¶

• DMLC_INTERFACE Using a particular network interface Value type: Name of interface Example: eth0 MXNet often chooses the first available network interface. But for machines with multiple interfaces, we can specify which network interface to use for data communication using this environment variable.
• PS_VERBOSE Logging communication Value type: 1 or 2 Default value: (empty)
  • PS_VERBOSE=1 logs connection information like the IPs and ports of all nodes
  • PS_VERBOSE=2 logs all data communication information

When the network is unreliable, messages being sent from one node to another might get lost. The training process can hang when a critical message is not successfully delivered. In such cases, an additional ACK can be sent for each message to track its delivery. This can be done by setting PS_RESEND and PS_RESEND_TIMEOUT

• PS_RESEND Retransmission for unreliable network Value type: 0(false) or 1(true) Default value: 0 Whether or not to enable retransmission of messages
• PS_RESEND_TIMEOUT Timeout for ACK to be received Value type: Integer (in milliseconds) Default value: 1000 If ACK is not received in PS_RESEND_TIMEOUT milliseconds, then the message will be resent.