2bit gradient compression

[rahul003]

Description

Implements 2bit gradient compression by quantizing each value in gradient array to 2bits using user specified threshold. Shows about 2x speedup on large models with components like fully connected layers, and LSTM layers.

@eric-haibin-lin @cjolivier01 @anirudh2290 @reminisce

Important files to review

GC : gc-inl.h, gc.cc
KVStore local: comm.h
KVStore dist : kvstore_dist.h, kvstore_dist_server.h
Documentation about gradient compression: kvstore.py

Checklist

Essentials

• Passed code style checking (make lint)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• For user-facing API changes, API doc string has been updated.
• To my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Gradient compression class
• Reduce operation in kvstore_local / comm.h
• Distributed kvstore changes at worker and server
• Tests for local kvstore, distributed kvstore with predefined and random data. The results have been compared with expected values by implementing this logic in python.
• API changes for Kvstore, Module and Trainer in python
• Addressed comments from last PR

Comments

Problem

When training large scale deep learning models especially with distributed training, communication becomes a bottleneck for networks whose computation is not high compared to the communication.

Approach

We can compress the gradients by considering only those elements that exceed a threshold. Only these elements are encoded and sent. The elements of the gradient that are near zero can safely be delayed by aggregating them in a residual array. When the updated residual with gradient of next iterations exceed the threshold, these values are sent. Effectively these values are updated at a lower frequency.
On the receiver's end we decompress the data and use the decompressed weights.
Specifically in this PR, 2bit quantization has been implemented.

Two bit quantization

Any positive value greater than or equal to the threshold is set to one value (say 11), any negative value whose absolute value is greater or equal to the threshold is set to second value (say 10), and others are set to third value (say 0). We need three values to represent data in this fashion and hence two bits. We understand this leads to one bit going waste, but that's an optimization to be done later. The error in quantization is accumulated as residual and carried over to the next iterations. This is added in the next iteration to the gradient before quantizing.
An example below with thresholds of -2.0 and 2.0.
This format leads to the reduction of gradient size by 1/16th.
Quantization at work

Format of compressed gradient

Eac element, represents upto 16 elements in the original array. For the example above, we get an element whose binary representation is
00 11 00 10 11 00 00 10 0000000000000000

Local kvstore

When using local kvstore, gradients compression only happens when using device communication. When gradients are pushed, before summing them up (Reduce), quantization and dequantization happen.
Example: Say we have 4 GPUs, and the gradients are being summed up on GPU0. Each device quantizes gradients, then sends quantized gradient to GPU0, which performs dequantization of this data before merging it with values from other GPUs. Note that here, there is no need to quantize gradients from GPU0 itself, but it is still being done so that there is no bias for the samples which were processed by GPU0.

Dist kvstore

When the set_gradient_compression method for kvstore is called, each worker sets those compress params and one worker sends these params to all servers. From then on, when before each value is pushed to the server, it is quantized. The server dequantizes the data and stores it as an array of the original size. When values are pulled from the server, it returns an array of the original size. The same happens when each server is handling shards of the data.

Usage

The reason I used a dictionary compress_params for the arguments was to ensure uniformity when we extend this to other quantization techniques. This is because each technique might take different type and number of parameters.

KVstore

kv = mx.kv.create('dist_sync')
kv.set_gradient_compression({'type':'2bit', 'threshold':0.5})

Module

mod = mx.mod.Module(net, compression_params={'type':'2bit', 'threshold':0.5})

Gluon Trainer

trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': .1}, 
                        compression_params={'type':'2bit', 'threshold':0.5})

Results

Summary
Shows about 2x speedup when models are large, have fully connected components, for distributed training. On local training, speedup is about 1.2x when there is no P2P communication.

1. For MLP with 4 fully connected layers of 1500 size and one fully connected layer of 3000 size

input dim	model size	speedup	batch size(on each gpu)
300	50MB	1.7x	256
300	50MB	1.4x	1024
150000		2x	64
270000	900MB	2x	128

For smaller models, the overhead of launching OMP threads is costing a bit, to get around it (if training using GPUs), setting OMP_NUM_THREADS=1 results in gradient compression is needed.

2. Shows speedup when communication is expensive. The above speedup was seen on g2.8x machines which have lower network bandwidth than p2.16x machines. p2.16x didn't see as much speedup.

3. Network types
   On models for imagenet input (input dim: 3,299,299), on g2.8x large, 15 node cluster, used all 4 gpus on each node

network type	speedup
LSTM, BiLSTM	about 1.25-1.5x
VGG11	1.8x
MLP, Alexnet	2x

4. Accuracy
   LSTM on PennTreeBank with 200dim 2 layers
   

MNIST on MLP

CIFAR with resnet

Accuracy starts off slow, but the network converges to similar accuracy.
Accuracies at a few epochs
epoch 101 :
2bit: 0.80645, none: 0.83572, difference: 0.029
epoch153
2bit: 0.841, none: 0.851, difference: 0.0108

CIFAR resnet with pretraining

Pre-training without gradient compression for some time(2 epochs), leads to better convergence
We see that in this case, we start off much closer and reaches similar accuracies earlier. In general, the graphs are much closer. Let's look at epoch 33, Earlier, without pretraining, 2bit compression had an accuracy degradation of 0.154 when compared to the case without gradient compression. Now, when both models start with a pretrained network which didn't use gradient compression, it has a degradation of only 0.04.

Reference (although compressed representation is different http://nikkostrom.com/publications/interspeech2015/strom_interspeech2015.pdf )

[piiswrong]

I don't think there should be a default value at all.

[piiswrong]

rename key compression to type

[piiswrong]

Not necessary. gc_ defaults to nullptr

[piiswrong]

Use scoped enum.
enum class CompressionType{
kNone,
kTwoBit
};

[piiswrong]

Use enum for this

[piiswrong]

These parsing should be done in backend with dmlc::Parameter

The frontend should pass strings of key value pairs.

[piiswrong]

API should be
MXKVStoreSetGradientCompression(KVStoreHandle handle, mx_uint num_params, const char **keys, const char **vals)
The values should be parsed in backend with dmlc::Parameter

[rahul003]

@piiswrong Updated to use scoped enums, and DMLC param

Wanted to add that tests are all in nightly because this affects either distributed kvstore dist_* or device both of which can't be tested in unittests

[eric-haibin-lin]

no support for un-setting gradient compression ? What happens if an user tries to unset it?

[rahul003]

if user uses kvstore.set_gradient_compression({'type':'none'} after setting it to 2bit, it throws an error because none can't be a type.
If users sets 2bit again with different threshold, then new threshold will be used from then on, but there might be a period in transition when gradients quantized with old threshold will be dequantized with new threshold, because of delay in sychronization.

[rahul003]

Does this look ready to be merged now?
The build had passed but status wasn't communicated to GitHub because of some Jenkins issues I guess. It will run again soon.

I've updated the results section with more details.

I'll hopefully be updating gradient compression with more optimizations before the v1.0 release. But it would be better if we merge this, so next PRs aren't this large. There's no known bug right now.