What to do with to(device) when using DataParallel?

scientia: <p>It’s not clear to me what to do with <code>.to(device)</code> when using  <code>DataParallel</code>.</p>
<p>Imagine something like</p>
<pre><code class="lang-auto">device = torch.device('cuda:2')
model = nn.DataParallel(model, device_ids=[2,3])
</code></pre>
<p>In my model I now have:</p>
<pre><code class="lang-auto">def forward(self, input):
    input = input.to(device)
</code></pre>
<p>Outside of my model I have</p>
<pre><code class="lang-auto">output = output.to(device)
torchmetrics.Accuracy(num_classes=1).to(device)
</code></pre>
<p>It raises following questions:</p>
<ul>
<li>Is <code>to(device)</code> now unnecessary? Should I remove it?</li>
<li>If not, which ID give I to the device?</li>
<li>Is the <code>to(device)</code> handled the same for each case (inside the forward function, outside of it)?</li>
</ul>

aazzolini: <aside class="quote no-group" data-username="scientia" data-post="1" data-topic="156649">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/b3f665/48.png" class="avatar"> scientia:</div>
<blockquote>
<p><code>DataParallel</code></p>
</blockquote>
</aside>
<p>For the input: You should remove input = input.to(device) from inside your forward function. Once you wrap your model with DataParallel, the wrapper will take care of splitting the inputs and placing them on the right device for you. forward will be called once for each device, so you shouldn’t move the input to a specific device.</p>
<p>As for the output: the documentation here <a href="https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html" class="inline-onebox" rel="noopener nofollow ugc">DataParallel — PyTorch 2.1 documentation</a>  mentions that by default output_device is set to devices[0] (probably your ‘cuda:0’ device). But you can actually pass the desired output device as well to the constructor of DataParallel.</p>
<p>Let me know if you have trouble with this.</p>

scientia: <p>Hmm, if I’m not supposed to have <code>.to(device)</code>, then I have trouble understanding how PyTorch will know which tensors will need to move to the GPU, esp. for my use case: <code>input</code> is actually a list in my model, and I repeatedly call a subnetwork inside the <code>forward</code> function (for each list entry, in each iteration I call <code>.to(device)</code>) before I use its output on my top-level network part (for which I also call <code>.to(device)</code>).</p>

aazzolini: <p>It’s a bit magical, I agree , but yes, that’s what it does. It goes into your input recursivelly (if it is a list it goes into each element of the list, etc, here: <a href="https://github.com/pytorch/pytorch/blob/557fbf4261d6517552b48c47be9aa9d289fa28d3/torch/nn/parallel/scatter_gather.py#L25-L36" class="inline-onebox" rel="noopener nofollow ugc">pytorch/scatter_gather.py at 557fbf4261d6517552b48c47be9aa9d289fa28d3 · pytorch/pytorch · GitHub</a></p>
<p>then for each Tensor it finds, it ultimatelly calls: <a href="https://github.com/pytorch/pytorch/blob/557fbf4261d6517552b48c47be9aa9d289fa28d3/torch/nn/parallel/comm.py#L152" class="inline-onebox" rel="noopener nofollow ugc">pytorch/comm.py at 557fbf4261d6517552b48c47be9aa9d289fa28d3 · pytorch/pytorch · GitHub</a>, which scatters the tensor across GPUs.</p>
<p>Let me know if, after trying it out, you have issues.</p>

scientia: <p>I tried to use it as a drop-in replacement, so the same setup as in the first post:</p>
<pre><code class="lang-auto">device = torch.device('cuda:4')
model = nn.DataParallel(model, device_ids=[4,5])
</code></pre>
<p>and as you suggested I removed the <code>to(device)</code> calls from inside the <code>forward</code> function – but I run into this error:</p>
<pre><code class="lang-auto">RuntimeError: module must have its parameters and buffers on device cuda:4 (device_ids[0]) but found one of them on device: cpu
</code></pre>
<p>So, after some research, it seems, you have to call <code>to(device)</code> on the model and input before passing it to <code>forward</code>, basically moving it to the primary GPU from where it’ll be scattered to other ones? That’s how I understand it right now. So after that, and moving the input tensors to the same device, it stops returning that error message.</p>
<pre><code class="lang-auto">model = nn.DataParallel(model, device_ids=[4,5]).to(device)
</code></pre>
<pre><code class="lang-auto">model([[x.to(device) for x in y] for y in inputs], lengths.to(device), g=False)
</code></pre>
<p>… and allows me to reach this point in my <code>forward</code> function, which previously had <code>to(device)</code> and now doesn’t, where I dynamically create a padding mask:</p>
<pre><code class="lang-auto">mask = torch.arange(max_len).expand(len(lengths), max_len) &lt; lengths.unsqueeze(1)
</code></pre>
<p>And returns this error:</p>
<pre><code class="lang-auto">RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:4 and cpu!
</code></pre>
<p>So I’m not sure, if I’m now totally on the wrong track but the last error at least suggest that I need some kind of <code>to(device)</code> inside <code>forward</code>, but of course it should be the corresponding device and not a fixed one.</p>

ptrblck: <aside class="quote no-group" data-username="scientia" data-post="5" data-topic="156649">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/b3f665/48.png" class="avatar"> scientia:</div>
<blockquote>
<p>So I’m not sure, if I’m now totally on the wrong track but the last error at least suggest that I need some kind of <code>to(device)</code> inside <code>forward</code>, but of course it should be the corresponding device and not a fixed one.</p>
</blockquote>
</aside>
<p>If you need to create new tensors in the <code>forward</code> method, push them to the corresponding device by using the <code>.device</code> attribute of any parameter or the input:</p>
<pre data-code-wrap="python"><code class="lang-python">mask = ...
mask = mask.to(x.device) # x is the input to forward
</code></pre>

scientia: <p>Moving the newly created tensor to the GPU seems to work. However, there’s a dimension mismatch.</p>
<p>It seems it doesn’t magically split my list of list into two batches of lists of lists, so it just passes the whole object to both GPUs. On the other hand, it seems to work for the <code>lengths</code> tensor causing a dimension mismatch at some point.</p>
<pre><code class="lang-auto">AssertionError: expecting key_padding_mask shape of (28, 49), but got torch.Size([14, 49])
</code></pre>

scientia: <p>Actually, looking more closely at the <code>scatter_map</code> function, it might be that in case of a list (or list of lists), it doesn’t split the list in half*, but each individual tensor, meaning that both GPUs get a list of the same size as the original one, but with fractured tensors? Is that right? And if it’s the case, is there any way to change that behaviour?</p>
<p>*Assuming two GPUs.</p>