Which GPU is used as KVStore device on multi-GPU machine?

[bill10]

Hi mxnet community, I have spent a whole day googling around but could find this information: when training on a multi-gpu machine with the gluon.Trainer and kvstore set to "device", which gpu is used as the kvstore? Or putting it differently, how to specify which gpu to use?

Based on my tests, it is not trivial because on a AWS p2.8x, gpu(0) is used as the kvstore but on a AWS p2.16x, gpu(10) is used. Can't find any documentation about the logics behind it. Any ideas are highly appreciated. Thanks!

[leezu]

While not a direct answer to your question, please consider using https://github.com/horovod/horovod/ to speed up the multi-gpu training. You will want to have a separate process for each GPU where you set CUDA_VISIBLE_DEVICES to ensure that the process only sees it's own GPU. Then you want to use horovod to communicate accross the processes / gpus. This will be significantly faster than mxnet kvstore as you reduce the Python overhead by using multiple Python processes.

[bill10]

@leezu thanks for the suggestion! Yeah, if I were to start a new project I would probably switch to horovod. However, this project is already code complete with only this bug. Or it is not really a bug but something annoying since I need to hardcode the kvstore device number. Anyway, thanks for the idea! If you have any reference about the mxnet code, please send them along and I will dig into them. Thanks again!

[szha]

The "device" kvstore attempts at instantiating the buffers for gradient aggregation evenly across all seen devices: https://github.com/apache/incubator-mxnet/blob/v1.x/src/kvstore/comm.h#L677-L712.

The problem is that the Gluon trainer only initiates it on the first context: https://github.com/apache/incubator-mxnet/blob/v1.x/python/mxnet/gluon/trainer.py#L158. As a result, all gradient aggregation only happens on gpu(0).

[bill10]

Thanks for the pointer! Looks like the initialization is from param_arrays[0] but not exactly gpu(0). That's probably why on p2.16x I saw the kvstore is on gpu(9) or gpu(10). I will check it again tomorrow and see if the context of param_arrays[0] matches what I saw. Thanks!

[bill10]

Just checked param_arrays[0] as in the line referenced above. It is always data from gpu(0). This does not agree with what I saw on P2.16x where gradient aggregation happened on gpu(10). Super confused.