How to apply DDP in an algorithm which uses model parameters outside the forward function

Yicheng-G: <p>Hi, I find it quite difficult to use DDP to train a model with an additional loss function outside the <code>forward</code> function.</p>
<hr>
<h3>Training Procedure</h3>
<p>The model (<strong>M</strong>, based on ProxylessNAS) has two sets of parameters,</p>
<ul>
<li>neural network weights <strong>W</strong>,</li>
<li>architecture parameters (operator weights) <strong>A</strong>,</li>
</ul>
<p>The steps to update <strong>A</strong> are,</p>
<ol>
<li>randomly sample a sub-network <strong>M<sub>sub</sub></strong>, with parameters <strong>W<sub>sub</sub></strong>, based on the probability matrix <strong>A</strong>
</li>
<li>
<code>loss1 = </code><strong>M<sub>sub</sub></strong><code>(data)</code>
</li>
<li>
<code>loss2</code> is directly calculated from <strong>A</strong>, e.g. <strong>Latency(A) = 3xA<sub>01</sub> + 2xA<sub>02</sub></strong>
</li>
<li><code>loss = f(loss1, loss2); loss.backward()</code></li>
<li>update <strong>A</strong>…</li>
</ol>
<hr>
<h3>The Problem</h3>
<p>I set <code>find_unused_parameters=True</code> and it raises an error during backward propagation.</p>
<pre><code class="lang-bash">RuntimeError: Expected to mark a variable ready only once. This error is caused by use of a module parameter outside the `forward` function. The return value of the `forward` function is inspected by the distributed data parallel wrapper to figure out if any of the module's parameters went unused. If this is the case, it knows they won't receive gradients in a backward pass. If any of those parameters are then used outside `forward`, this error condition is triggered. You can disable unused parameter detection by passing the keyword argument `find_unused_parameters=False` to `torch.nn.parallel.DistributedDataParallel`.
</code></pre>
<p>The problem is since only a part of the model (<strong>M<sub>sub</sub></strong>) is used in each iteration,  DDP won’t get gradients of those parameters not belong to <strong>W<sub>sub</sub></strong>. If I set <code>find_unused_parameters=False</code>, it will crash in next forward pass.</p>
<p>Does anyone have any idea to solve the problem?</p>

rvarm1: <p>In this case, I’m assuming that A is a part of the network M, although during backwards, since according to DDP, A went unused, it is marked as ready for reduction, but then we attempt to re-mark A’s parameters when it gets grads in the backwards pass.</p>
<p>Is it possible to create a wrapper network that wraps steps 2 and 3 into a single module and returns both loss1 and loss2/the values needed to compute it? That seems like it would avoid the double mark issue.</p>
<p>Also, the version of PyTorch you are using and a script to repro the issue would also be valuable.</p>

Yicheng-G: <blockquote>
<p>In this case, I’m assuming that A is a part of the network M, although during backwards, since according to DDP, A went unused, it is marked as ready for reduction, but then we attempt to re-mark A’s parameters when it gets grads in the backwards pass.</p>
</blockquote>
<p>Yes, this is the cause. Thanks for the explanation.</p>
<blockquote>
<p>Is it possible to create a wrapper network that wraps steps 2 and 3 into a single module and returns both loss1 and loss2/the values needed to compute it? That seems like it would avoid the double mark issue.</p>
</blockquote>
<p>That’s a good idea. Although it means a lot of work <img src="https://discuss.pytorch.org/images/emoji/apple/rofl.png?v=12" title=":rofl:" class="emoji" alt=":rofl:" loading="lazy" width="20" height="20">.</p>
<p>I’m using pytorch 1.4. The code is a little bit messy, I’ll try to make a demo if above suggestion not works.</p>