Does DDP with torchrun need torch.cuda.set_device(device)?

coni: <p>I have a training scripts that runs on single node, multiple GPUs, implemented following <a href="https://pytorch.org/tutorials/beginner/ddp_series_fault_tolerance.html" rel="noopener nofollow ugc">PyTorch DDP tutorial</a>. I run the script with <code>torchrun --standalone --nproc_per_node=8 main.py</code></p>
<p>According to the <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html" rel="noopener nofollow ugc">docs</a>:</p>
<blockquote>
<p>To use <code>DistributedDataParallel</code> on a host with N GPUs, you should spawn up <code>N</code> processes, ensuring that each process exclusively works on a single GPU from 0 to N-1. This can be done by either setting <code>CUDA_VISIBLE_DEVICES</code> for every process or by calling:<br>
<code>torch.cuda.set_device(i)</code> where i is from 0 to N-1.</p>
</blockquote>
<p>However, this is not mentioned at all in the PyTorch tutorial on DDP.<br>
Is this really needed? What does <code>set_device()</code> do in practice? Which are the implications of not calling it? Should I set <code>CUDA_VISIBLE_DEVICES</code> instead?</p>
<p>I am trying to understand more on this topic, thank you in advance!</p>

suraj.pt: <p><a class="mention" href="/u/coni">@coni</a> Tutorial author here, thanks for bringing up this discrepancy! Although scripts can work without <code>set_device(i)</code>, you should include it on your training script. Doing so changes the default GPU from <code>cuda:0</code> to <code>cuda:i</code>. Not doing this in some cases can result in your training job hanging, or too much memory utilization on <code>cuda:0</code>.</p>
<p>I will update the tutorial to explicitly call this out.</p>

coni: <p>Thanks you for the prompt response and for the tutorial!</p>

coni: <p><a class="mention" href="/u/suraj.pt">@suraj.pt</a> Thanks for the response, I have a follow-up.<br>
What if I wanted to set env variable <code>CUDA_VISIBLE_DEVICES</code>, instead of calling <code>set_device</code>?</p>
<p>I have tried to set in the command line, by calling</p>
<blockquote>
<p>CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --standalone --nproc_per_node=8 main.py</p>
</blockquote>
<p>However when checking the current device inside <code>main.py</code> with <code>print(torch.cuda.current_device())</code>,  I always get <code>0</code>, is this behavior correct? Does <code>cuda:0</code> in different processes map to a different physical device?</p>

suraj.pt: <p><code>cuda:0</code> across all the processes maps to the same GPU. You are always getting <code>0</code> because <code>cuda:0</code> is the default device. Unless you run <code>set_device</code> (or something equivalent like <code>torch.device</code>) this will not change.</p>
<p>See <a href="https://github.com/pytorch/examples/pull/1142/files" class="inline-onebox">Adds torch.cuda.set_device calls to DDP examples by subramen · Pull Request #1142 · pytorch/examples · GitHub</a>.</p>
<p>You can set <code>CUDA_VISIBLE_DEVICES=i</code> from each <code>i</code>th process, which will have the same effect as <code>set_device</code>. Replace <code>torch.cuda.set_device(rank)</code> with <code>os.environ["CUDA_VISIBLE_DEVICES"]=rank</code> and it will set the visible devices for that specific process.</p>

coni: <p>Setting <code>os.environ["CUDA_VISIBLE_DEVICES"]=os.environ["LOCAL_RANK"]</code> still results in <code>torch.cuda.current_device()</code> equal to <code>0</code>.</p>
<p>However, from <a href="https://discuss.pytorch.org/t/should-local-rank-be-equal-to-torch-cuda-current-device/150873">this topic</a> this behavior seems correct. It seems that by setting <code>CUDA_VISIBLE_DEVICES</code> directly “<em>each process will only see one physical GPU that corresponds to its <code>local_rank</code> , i.e., <code>cuda:0</code> in different processes will map to a different physical device</em>”.</p>
<p>Is this correct?</p>

suraj.pt: <blockquote>
<p><em>each process will only see one physical GPU</em></p>
</blockquote>
<p>Given the naming of the env var, that does make sense… I am not sure if <code>set_devices</code> works differently than setting <code>CUDA_VISIBLE_DEVICES</code></p>
<p>Paging <a class="mention" href="/u/rvarm1">@rvarm1</a> / <a class="mention" href="/u/mrshenli">@mrshenli</a> to chime in here</p>

coni: <p>Adding something more to this, <code>torch.device()</code> produces the same behavior of manually setting CUDA_VISIBLE_DEVICES.</p>
<p>Summarizing:</p>
<pre><code class="lang-auto">torch.set_device(rank) # results in torch.cuda.current_device() equal to rank
torch.device(rank) # results in torch.cuda.current_device() equal to 0
os.environ["CUDA_VISIBLE_DEVICES"]=rank # results in torch.cuda.current_device() equal to 0
</code></pre>

ptrblck: <aside class="quote no-group" data-username="suraj.pt" data-post="7" data-topic="178723">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/suraj.pt/48/53186_2.png" class="avatar"> suraj.pt:</div>
<blockquote>
<p>Given the naming of the env var, that does make sense… I am not sure if <code>set_devices</code> works differently than setting <code>CUDA_VISIBLE_DEVICES</code></p>
</blockquote>
</aside>
<p>Setting <code>CUDA_VISIBLE_DEVICES=id</code> masks the available GPUs and allows the process to only use the visible device. This visible device will be mapped to <code>cuda:0</code> inside your Python script and no other devices can be accessed thus <code>torch.cuda.device_count()</code> will also return <code>1</code>. (Multiple visible devices will map to <code>cuda:0</code>, <code>cuda:1</code>, etc. in the same order specified in the env variable).<br>
Using <code>torch.cuda.set_device</code> will set the specified device as the current device and will execute all operation on it if no <code>cuda:ID</code> is specified explicitly. All other visible devices are still accessible in this process and your script.</p>

suraj.pt: <p><a class="mention" href="/u/coni">@coni</a> your findings match with <a class="mention" href="/u/ptrblck">@ptrblck</a>’s explanation here.</p>
<p>Looks like <code>CUDA_VISIBLE_DEVICES</code> / <code>torch.device()</code> are actually more straightforward to use for typical DDP applications.</p>
<p>The only time you would  <em>need</em> to use <code>torch.set_device()</code> instead of the above is when your script needs to access multiple GPUs from the same process - not a typical need when using torchrun.<br>
That’s probably why the <a href="https://pytorch.org/docs/stable/generated/torch.cuda.set_device.html?highlight=set_device#torch.cuda.set_device">docs</a> also mention using <code>CUDA_VISIBLE_DEVICES</code> / <code>torch.device()</code> instead of  <code>torch.set_device()</code></p>

Nitish_Kumar1: <p>Hello Suraj,</p>
<p>Thanks for the tutorial.</p>
<p>One question I ahve when we are using the multinode and in that when we are running using torchrun.<br>
Let say I have two node each having gpu id 0,1 2 3 but on first node I want to use gpu id 0 and 1 and on node 2 I want to use gpu id 1,2 and 3 then how can we set these ids?</p>
<p>Thanks in advance.<br>
Nitish</p>