Sub-path training using DDP may lead incorrect optimizing!

username: VisionAdam
cooked: <p>In sub-path training, such as NAS，there are two path A and B need to be optimized. for example, the optimizing order is B-&gt;A:<br>
1.when using a single-gpu, everything works well if we set require_grad=False for unselected path，<br>
2.but when using DistributedDataParallel, we need set find_unused_parameters=True in this condition，and because of  grad sync，for example, when optimizes path A in step2, the parameters‘ grad in path B will be 0 rather then None<br>
Here comes the problem: optimizer chooses to update params which grad is not None. If we set momentum in SGD or something else in other  optimizer, when optimizing path A in step 2, because we have already optimized path B once，so parameters in B have momentum_buffer, although path B’s grad is 0,  it will also be updated by the buffer which is incorrect!<br>
Even we set require_grad=False for path B, but once path B is optimized, its grad will be 0 rather then None.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/c/fc064781b536fa89ca9fd695aee7e0a762c9b006.png" data-download-href="https://discuss.pytorch.org/uploads/default/fc064781b536fa89ca9fd695aee7e0a762c9b006" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fc064781b536fa89ca9fd695aee7e0a762c9b006_2_690x316.png" alt="image" data-base62-sha1="zXvM6IlAIJrHJuG2zwZhhicspGm" width="690" height="316" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fc064781b536fa89ca9fd695aee7e0a762c9b006_2_690x316.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fc064781b536fa89ca9fd695aee7e0a762c9b006_2_1035x474.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fc064781b536fa89ca9fd695aee7e0a762c9b006_2_1380x632.png 2x" data-dominant-color="33302F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1712×786 65.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: wanchaol
cooked: <p>Thanks for posting <a class="mention" href="/u/visionadam">@VisionAdam</a> Do we know why the grad in path B becomes 0 even if we optimizing path A? I recalled DDP will just don’t touch the gradient field if it’s not being used in forward. Can you share more details why we are observing those gradients be 0?</p>
<p>Also, did you try setting the gradients of Path B explicitly to None and see if that fix your issue?</p>

username: rvarm1
cooked: <p>If <a class="mention" href="/u/wanchaol">@wanchaol</a>’s suggests don’t resolve the issue, feel free to post this as a bug over at <a href="https://github.com/pytorch/pytorch/issues" class="inline-onebox" rel="noopener nofollow ugc">Issues · pytorch/pytorch · GitHub</a> ideally with a repro!</p>

username: VisionAdam
cooked: <p>if we never optimize path B, the grad will be None, but if path B optimized once, the grad will be 0 in future time.<br>
It can be fixed by set grad to None while the grad is all zero</p>