Will `backward()` consider all the models across various GPUs with `nn.DataParallel`?

smmislam: <p>I have a model like the following</p>
<pre><code class="lang-auto">class MyModel(nn.Module):
    ...
    def forward(self, x):
        ...
        return y

    def compute_loss(self, y, t):
        ...
        return loss

    def update(self, loss):
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
</code></pre>
<p>I am trying to run this model on two GPUs using <code>DataParallel</code>. And my <code>main()</code> looks something like this</p>
<pre><code class="lang-auto">DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = nn.DataParallel(MyModel).to(DEVICE)
...
for xb, yb in train_dl:
    y = model(xb) # Data being fed to two GPUs, I'm happy :)
    loss = model.module.compute_loss(y, yb)  # This is the issue
    model.module.update(loss) # This too!
</code></pre>
<p>I am wondering, <strong>will this consider the two different versions of the model living in two different GPUs during weight update?</strong></p>
<p>My code is running on both the GPUs (checked with <code>nvidia-smi</code>), but I am having doubt that even though during <code>forward()</code> both the GPUs are being used, <code>compute_loss</code> and <code>update</code> are only processing loss/gradients for the model living at <code>cuda:0</code>, because I’m calling <code>compute_loss</code> and <code>update</code> via <code>module</code>.</p>
<p>Or it could be that the <code>loss</code> is linked to both the GPU’s gradients through the computation graph since <code>model(xb)</code> returned combined <code>y</code> from both the GPUs and because of that <code>loss.backward()</code> would run through both of them.</p>
<p>Could you please shed some lights on my doubt?<br>
Thank you so much! <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=12" title=":grin:" class="emoji" alt=":grin:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>The loss will be computed on the default device, but the backward pass should be executed on both devices again as the loss calculation should still be attached for the computation graph created in the forward on both devices.</p>

smmislam: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you for your comment!</p>
<p>Meanwhile I was able to find another very useful <a href="https://discuss.pytorch.org/t/how-pytorchs-parallel-method-and-distributed-method-works/30349/2">thread</a> where a beautiful illustration of how <code>DataParallel</code> actually works in the background is shared.</p>
<p>I think that the code that I shared above, might not work properly, because I declared the optimizer as part of the model. And, since I’m calling the optimizer through <code>module</code>, it will only update the default GPU’s model weights. Even if the errors manages to flow back to both the models because of linked computation graph, I do not see any routine to merge the gradients as one and make one global update.</p>
<p>However, if I instantiate optimizer independently and call it on <code>model.parameters</code> from the <code>main()</code> step 5 and 6 (from the illustration) should run as intended.</p>
<p><a class="mention" href="/u/rasbt">@rasbt</a> Thank you so much for the illustration. Do you think my above remark is correct?</p>

ptrblck: <p>I don’t think the description is completely correct for <code>nn.DataParallel</code> as you would be handling a single model only and the gradients should be already reduced to the default model on the default device.<br>
However, we generally don’t recommend <code>nn.DataParallel</code> as it’s slower than <code>DistributedDataParallel</code> so I also don’t know who the “internal” optimizer might behave.<br>
In any case, since you have concerns I would recommend to run a quick test using defined data to check the behavior. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

smmislam: <p>I could not really see which specific command would cause this reduction of gradients. Unless there is some mechanism inside the <code>optimizer</code> to check if there are multiple GPUs involved and hence reduce the gradients automatically.</p>
<p>But, I would like to do the quick test you mentioned. What do you mean by <strong>defined data</strong>? And how to do such tests? Would really appreciate if you could point me somewhere with a bit more explanation.</p>
<p>Apologies for the noob questions. <img src="https://discuss.pytorch.org/images/emoji/apple/frowning_face.png?v=12" title=":frowning_face:" class="emoji" alt=":frowning_face:" loading="lazy" width="20" height="20"></p>

ptrblck: <p><a href="https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255">This blog post</a> as well as the post you’ve linked in your previous method show where the gradients are reduced to the default device.</p>