Horovod GPU memory imbalance on parallel/distributed training

[feevos]

Dear all,

during distributed training with Horovod on our local HPC facilities (~100 nodes with 4xP100 GPUs each), we observed different memory footprint differences between GPU in local rank 0, and the other GPUs. This is documented in horovod issue 2625 here.

This solution fixed it for me, leaving it for reference for other users:

from mpi4py import MPI
rank = MPI.COMM_WORLD.Get_rank()
num_devices=4 # 4 is the number of GPUs available per node. 
local_rank = rank % num_devices  

import os
os.environ['CUDA_VISIBLE_DEVICES']=str(local_rank)

import mxnet as mx
mx.npx.set_np()

# ..... more code .... 
 
ctx = mx.gpu(0)

Regards,
Foivos

[leezu]

Thank you @feevos. Do you think Horovod should set this variable automatically during startup via horovodrun?

[feevos]

Hi @leezu thank you for your reply. I am afraid I don't know, I am out of my depth here. It is for sure helping with memory balance in training memory intensive models. In my case I was having ~3GB memory difference between local rank 0, and each of the local ranks 1,2,3. That's significant in a 16GB GPU. This was for a very heavy model (semantic segmentation), with batch size of 2 (per GPU).

Regards,
Foivos

[leezu]

No worries. I think horovod should configure this variables. But after second thought I think they cannot do it for backwards compatibility reason (given you now need to use mx.gpu(0) instead of mx.gpu(local_rank)).