Huge loss with DataParallel

username: neale
cooked: <p>I’m training a pretty standard WGAN-GP on MNIST, and I’m trying out much larger batch sizes since that seems to be the standard wisdom now. When I parallelize across multiple GPUs I get enormous losses compared to using just one GPU. If I initialize my networks with <code>nn.DataParallel(model, [0])</code> then I get pretty normal functionality:</p>
<p>ITER  0: D cost 1.200, G cost 2.729<br>
ITER  200: D cost -3.931, G cost 2.298</p>
<p>But if I use <code>nn.DataParallel(model, [0, 1, 2])</code>  to run across more GPUs, I get absurd numbers:</p>
<p>ITER  200: D cost 112856899584.0, G cost 456269.437</p>
<p>I’ve used dataparallel successfully before with classifiers and the tutorial DCGAN, so I have no idea what the issue is here. I’m not parallelizing the loss or doing anything different than initializing the networks in this way. Are there some caveats here (like with FP16) that I’m not aware of with DataParallel?</p>
<p>I can post a full code example, but its a hundred lines or so and I’d rather not start off the post like that.</p>

username: pietern
cooked: <p>The DataParallel module is pretty straightforward: it splits the input up into N chunks (typically across the first dimension), runs the same forward pass on N replicas of your model, and gathers the output back into a single tensor (across the same dimension as the input was split). Gradients are always accumulated in the source model (not the replicas).</p>
<p>It looks like updates to buffers in these replicas don’t propagate back to the source model, as the model replicas are tossed after every forward pass. Perhaps this is a starting point for your investigation?</p>

username: neale
cooked: <p><a class="mention" href="/u/pietern">@pietern</a> even if the gradients didn’t accumulate properly, it should at worst look as if I only have one GPU. I’ve verified that everything works well with up to 4 GPUs with pytorch <code>0.4.1</code>. But I can’t seem to nail down the actual problem in <code>1.0</code></p>
<p>I’m not even sure how to diagnose this, but I’ve been able to replicate this behavior with dataparallel on 2 other popular WGAN repos from github.</p>

username: Hung_Nguyen
cooked: <p>I had the same problem with with wp-gan. For some reasons, the gradient penalty increases quickly until it reaches inf.</p>
<p>Also, this happens for other gan too, e.g Self supervised gan. The behavior of loss function is also completely different when training with multi gpus.</p>

username: neale
cooked: <p><a class="mention" href="/u/hung_nguyen">@Hung_Nguyen</a> Right.<br>
Remove the gradient penalty and the loss should still endlessly increase.<br>
The official(?) Wasserstein GAN code doesn’t suffer from this weird behavior with parallelization, so that could be a starting point. It uses <code>nn.parallel.data_parallel</code> - the functional version of <code>nn.DataParallel</code>, but I don’t know if there’s an interesting difference there.<br>
Converting any WGAN-GP repo into a regular WGAN mitigates the behavior somewhat, but weight clipping has its drawbacks.</p>

username: aluo-x
cooked: <p>Currently also experiencing this issue. The gradient penalty eventually goes to the millions before blowing up completely. <a class="mention" href="/u/neale">@neale</a>, I have tried to reproduce this issue with various popular WGAN-GP repos as well, and they also suffer from this.</p>
<p>Pytorch 1.0.1 on CUDA 10.1.</p>

username: pietern
cooked: <p><a class="mention" href="/u/neale">@neale</a> With 0.4 working and 1.0 not working this is clearly a regression. But we haven’t significantly modified (if at all) the data parallel wrapper between these versions. Can you check if the regression happened in 1.0.0 or 1.0.1? I have created <a href="https://github.com/pytorch/pytorch/issues/19024" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/19024</a> to track this.</p>

username: pietern
cooked: <p><a class="mention" href="/u/neale">@neale</a> Could you also try reproducing with the nightly build? There have been some changes recently related to CUDA stream synchronization that may have fixed this, per <a class="mention" href="/u/mrshenli">@mrshenli</a>.</p>

username: Nurble
cooked: <p>Running into this problem as well, using a WGAN-GP and it works perfectly on 1 GPU but the loss explodes when running on multiple GPUs.</p>
<p>Using CUDA 9.2, PyTorch 1.0.1. Working on installing the nightly build to see if there is any difference.</p>

username: neale
cooked: <p><a class="mention" href="/u/pietern">@pietern</a> Sorry this took quite some time.<br>
I can confirm that the issue persists into version 1.1</p>

username: aluo-x
cooked: <p>It seems like the issue is <a href="https://github.com/pytorch/pytorch/issues/16433" rel="nofollow noopener">#16433</a>.</p>
<p>A workaround would be to calculate the gradient-penalty directly (without calling a function to do so) and calling backward in the same scope.</p>
<p>For example the following code will explode on CUDA with multi-gpu:</p>
<pre><code class="lang-auto">gp = calc_grad_penalty(network, real_target, fake_target)
gp.backwards(retain_graph=True)
</code></pre>
<p>While the following does not:</p>
<pre><code class="lang-auto">### gp += torch.autograd.grad()
### etc. etc. code to calcualte GP
gp.backwards(retain_graph=True)
</code></pre>

username: pietern
cooked: <p>Thanks, <a class="mention" href="/u/aluo-x">@aluo-x</a>.</p>
<p>It is also the same as <a href="https://github.com/pytorch/pytorch/issues/16532" rel="nofollow noopener">#16532</a> and there has been an attempt at a fix. The problem lies somewhere deep in the guts of autograd. This has surfaced a couple of times and there should be a fix soon (and it should be included in the next stable release).</p>

username: kiddyboots216
cooked: <p><a class="mention" href="/u/pietern">@pietern</a> you mentioned that gradient accumulation only occurs in the source model, but according to <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py</a> “In the forward<br>
pass, the module is replicated on each device, and each replica handles a<br>
portion of the input. During the backwards pass, gradients from each replica<br>
are summed into the original module.” this implies that gradients are accumulated in the leaf nodes of each of the replicas.</p>

username: pietern
cooked: <p>FYI, excellent debugging from <a class="mention" href="/u/mrshenli">@mrshenli</a> and <a class="mention" href="/u/ezyang">@ezyang</a> deep in the guts of autograd led to <a href="https://github.com/pytorch/pytorch/pull/22983" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/22983</a> and this was merged yesterday. Please give the latest nightly builds a try to see if fixed the issue.</p>