CUDA issue with DDP: CUDA out of Memory error (rank 0)

ulee: <p>I am experiencing issues when using Vision transformer based model with DistributedDataParallel (DDP). While everything works fine with just DP, switching to DDP causes the following problems.</p>
<h1><a name="problem-1-cuda-out-of-memory-on-rank-0-gpu-1" class="anchor" href="#problem-1-cuda-out-of-memory-on-rank-0-gpu-1"></a>Problem 1. CUDA out of memory on rank 0 GPU</h1>
<p>This issues doesn’t occur with DP. Despite setting the device to rank and calling torch.cuda.empty_cache() before loading the model, memory allocation on rank 0 leads to CUDA out of memory.</p>
<pre><code class="lang-auto">    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    torch.cuda.set_device(rank)
    torch.cuda.empty_cache()
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    encoder = encoder.to(rank)
    mlp = mlp.to(rank)
    encoder = DDP(encoder, device_ids=[rank])
    mlp = DDP(mlp, device_ids=[rank])
...
        for bat_idx, (x, label) in enumerate(dataloader):
            x, label = x.to(rank), label.to(rank).long()
            emb = encoder(x)

  File "", line 282, in
    emb = encoder(x)
  File "/home/user/miniconda3/envs/deploy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/miniconda3/envs/deploy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/miniconda3/envs/deploy/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1593, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 430.00 MiB. GPU
</code></pre>
<h1><a name="problem-2-cuda-error-invalid-device-ordinal-2" class="anchor" href="#problem-2-cuda-error-invalid-device-ordinal-2"></a>Problem 2: CUDA Error: Invalid Device ordinal</h1>
<p>When setting CUDA_VISIBLE_DEVICES=“2,3,4,5” , I encountered ‘CUDA error: invalid device ordinal’ when calling ‘torch.cuda.set_device(rank)’. This error does not occur if I don’t set ‘CUDA_VISIBLE_DEVICES’.</p>
<pre><code class="lang-auto"> File "", line 212, in
    torch.cuda.set_device(rank)
  File "/home/user/miniconda3/envs/deploy/lib/python3.9/site-packages/torch/cuda/__init__.py", line 399, in set_device
    torch._C._cuda_setDevice(device)
RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
</code></pre>
<p>Any suggestions to address these issues would be greatly appreciated.<br>
Thank you!</p>

ptrblck: <p>Check the memory usage via <code>nvidia-smi</code> to see if e.g. multiple processes are running on the default device.</p>
<p>For the second issue: narrow down which line of code fails and make sure each process uses only its corresponding device.</p>

ulee: <p>For the first question, it was due to a misunderstanding on my part.<br>
I believed that if I set batch_size=256 for 4 GPUs, then each GPU would be given 64 samples, similar to how DP scatters the data.<br>
However, in DDP, the same batch size (256) was given to each GPU.</p>