Using spectral_norm with DistributedDataParallel makes backward() fail

username: bornabesic
cooked: <p>Hi!</p>
<p>My training works if I use mutiple GPUs on a single machine (i.e.<code> DataParallel</code>). However, if I try to train on two machines (by using <code>DistributedDataParallel</code>), I get the following error on <code>.backward()</code>:</p>
<pre><code class="lang-auto">one of the variables needed for gradient computation has been modified by
an inplace operation: [torch.cuda.FloatTensor [256, 256, 5, 5]] is at version 2;
expected version 1 instead. Hint: the backtrace further above shows the
operation that failed to compute its gradient. The variable in question was 
changed in there or anywhere later. Good luck!
</code></pre>
<p><code>torch.autograd.set_detect_anomaly(True)</code> points me to <code>spectral_norm</code>'s code that updates the weight:</p>
<pre><code class="lang-auto">  File ".../python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__                                                                                                                  
    hook(self, input)                                                                                                                                                                                                                        
  File ".../python3.6/site-packages/torch/nn/utils/spectral_norm.py", line 99, in __call__                                                                                                              
    setattr(module, self.name, self.compute_weight(module, do_power_iteration=module.training))                                                                                                                                              
  File ".../python3.6/site-packages/torch/nn/utils/spectral_norm.py", line 86, in compute_weight                                                                                                        
    weight = weight / sigma 
</code></pre>
<p>This definitely does not look like an inplace operation.<br>
The same error occurs even if I use  <code>DistributedDataParallel</code> for a single machine.</p>
<p>Any suggestions or ideas are more than welcome. Thanks in advance.</p>
<h3>Versions</h3>
<p>PyTorch: 1.1.0<br>
CUDA: 9.0.176</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/bornabesic">@bornabesic</a>,</p>
<ol>
<li>Can you try if setting <code>broadcast_buffers=False</code> in <code>DistributedDataParallel</code> constructor works for you?</li>
<li>If not, can you try PyTorch v1.4?</li>
<li>If it still does not work, could you please provide code for minimum repro? Thanks!</li>
</ol>

username: bornabesic
cooked: <p><a class="mention" href="/u/mrshenli">@mrshenli</a></p>
<p>The problem happens regardless of the value of <code>broadcast_buffers</code>.</p>
<p>I managed to narrow down my search for the source of the problem.<br>
The error occurs if I use <strong>multiple GPUs per machine</strong> AND <strong>multiple forward passes</strong> of the module before <code>backward()</code>.<br>
Otherwise, it works just fine.</p>

username: mrshenli
cooked: <blockquote>
<p><strong>multiple forward passes</strong>  of the module before  <code>backward()</code> .</p>
</blockquote>
<p>Do you mean running multiple forward pass on the same DDP instance before launching the backward pass? If so, it is expected to hit errors due to <a href="https://github.com/pytorch/pytorch/blob/df8d6eeb19423848b20cd727bc4a728337b73829/torch/nn/parallel/distributed.py#L454-L464" rel="noopener nofollow ugc"><code>prepare_for_backward</code></a>. But I would expect it throws a different error though. A work around for this is to wrap your multiple forward pass into one <code>YourModule.forward</code> function, and then use DDP to wrap <code>YourModule</code>.</p>