How to use `torch.nn.parallel.DistributedDataParallel` and `torch.utils.checkpoint` together

LicharYuan: <p>Hi, I have some problems using <code>torch.nn.parallel.DistributedDataParallel</code> (DDP) and <code>torch.utils.checkpoint</code> together.</p>
<p>It is Ok, if I set <code>find_unused_parameters=False</code> in DDP. The dilemma is that my network is a dynamic CNN, which will not forward the whole model during training, which means I have to set the <code>find_unused_parameters=True</code>… And, if I don’t use the <code>torch.utils.checkpoint</code>, my network is too large to run, leading to the OOM problem…</p>
<p>Therefore, what should I do to meet my demands?</p>
<p>There are some links to this question, but they not solve my problem.</p>
<ol>
<li><a href="https://github.com/pytorch/pytorch/issues/43135" class="inline-onebox" rel="noopener nofollow ugc">How to use torch.utils.checkpoint and DistributedDataParallel together · Issue #43135 · pytorch/pytorch · GitHub</a></li>
<li><a href="https://github.com/pytorch/pytorch/issues/24005" class="inline-onebox" rel="noopener nofollow ugc">Using `torch.utils.checkpoint.checkpoint_sequential` and `torch.autograd.grad` breaks when used in combination with `DistributedDataParallel` · Issue #24005 · pytorch/pytorch · GitHub</a></li>
</ol>
<p>Part of error report:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/f/9f3a15a8ad35a17798ca4b6bbe3dfcfe69488f44.png" data-download-href="https://discuss.pytorch.org/uploads/default/9f3a15a8ad35a17798ca4b6bbe3dfcfe69488f44" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f3a15a8ad35a17798ca4b6bbe3dfcfe69488f44_2_690x132.png" alt="image" data-base62-sha1="mIAmxgKgSPlR8aAocTvBFnwlqrq" width="690" height="132" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f3a15a8ad35a17798ca4b6bbe3dfcfe69488f44_2_690x132.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f3a15a8ad35a17798ca4b6bbe3dfcfe69488f44_2_1035x198.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f3a15a8ad35a17798ca4b6bbe3dfcfe69488f44_2_1380x264.png 2x" data-dominant-color="121212"><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">1875×360 114 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>
<p>Thanks for all the suggestions!!!</p>

mrshenli: <p>DDP does not work with <code>torch.utils.checkpoint</code> yet. One work around is to run forward-backward on the local model, and then manually run <code>all_reduce</code> to synchronize gradients after the backward pass.</p>

LicharYuan: <p>It’s ok in <code>find_unused_parameters=False</code>, so I guess maybe manually define <code>checkpoint</code> can solve my problem?</p>
<p>Thus, I modify the checkpoint to save the variable(indicating which part to forward), and get output again in backward with this variable.</p>
<p>But it makes new problems (Not Enough Values to unpack in backward)…</p>

mrshenli: <aside class="quote no-group" data-username="LicharYuan" data-post="3" data-topic="96338">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/licharyuan/48/28857_2.png" class="avatar"> LicharYuan:</div>
<blockquote>
<p>It’s ok in <code>find_unused_parameters=False</code> , so I guess maybe manually define <code>checkpoint</code> can solve my problem?</p>
</blockquote>
</aside>
<p>The reason <code>find_unused_parameters=True</code> does not work is because, DDP will try to traverse the autograd graph from output at the end of the forward pass when <code>find_unused_parameters</code> is set to True. However, with checkpoint, the autograd graphs are reconstructed during the backward pass, and hence it is not available when DDP tries to traverse it, which will make DDP think those unreachable parameters are not used in the forward pass (although they are just hidden by checkpoint). When setting <code>find_unused_parameters=False</code>, DDP will skip the traverse, and expect that all parameters are used and autograd engine will compute grad for each parameter <strong>exactly once</strong>.</p>
<blockquote>
<p>But it makes new problems (Not Enough Values to unpack in backward)…</p>
</blockquote>
<p>Could you please elaborate more on this problem?</p>

LicharYuan: <aside class="quote no-group" data-username="mrshenli" data-post="4" data-topic="96338">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mrshenli/48/12220_2.png" class="avatar"> mrshenli:</div>
<blockquote>
<p>The reason <code>find_unused_parameters=True</code> does not work is because, DDP will try to traverse the autograd graph from output at the end of the forward pass when <code>find_unused_parameters</code> is set to True. However, with checkpoint, the autograd graphs are reconstructed during the backward pass, and hence it is not available when DDP tries to traverse it, which will make DDP think those unreachable parameters are not used in the forward pass (although they are just hidden by checkpoint). When setting <code>find_unused_parameters=False</code> , DDP will skip the traverse, and expect that all parameters are used and autograd engine will compute grad for each parameter <strong>exactly once</strong> .</p>
</blockquote>
</aside>
<p>If so, then the new problem is caused by my codes. I will try to figure it out .</p>
<p>I understand the effect of <code>find_unused_paramerts</code>  now. Therefore, even my network has some unused parameters, but in each forward, each GPU updates the same unused parameters. In such a situation, I also can set  <code>find_unused_paramerts=False</code> all right ?</p>

LicharYuan: <p>I fix my problem like the FP16Optimizer in <a href="https://github.com/open-mmlab/mmdetection/blob/927d71a98ce2446193c626d26268df0c29922fb0/mmdet/core/fp16/hooks.py#L11" rel="noopener nofollow ugc">mmdetection</a>, which is similar to the Apex <code>delay all reduce</code>.</p>
<aside class="quote no-group" data-username="mrshenli" data-post="2" data-topic="96338">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mrshenli/48/12220_2.png" class="avatar"> mrshenli:</div>
<blockquote>
<p>One work around is to run forward-backward on the local model, and then manually run <code>all_reduce</code> to synchronize gradients after the backward pass.</p>
</blockquote>
</aside>
<p>And I realize the solution is just you said before… Thank you very much!!!</p>
<p>But it seems work no matter what value of <code>find_unused_parameters</code>. I wonder why it work when <code>find_unused_parameters=True</code>, the traverse should failed logically.</p>

mrshenli: <blockquote>
<p>But it seems work no matter what value of  <code>find_unused_parameters</code> . I wonder why it work when  <code>find_unused_parameters=True</code> , the traverse should failed logically.</p>
</blockquote>
<p>IDK, I would assume this would lead to autograd hook fire on a ready parameter, which should trigger an error in DDP. BTW, when you manually run allreduce, DDP is no longer necessary. Any reason for still wrapping the model with DDP?</p>

LicharYuan: <p>You are right. Actually, there is no need to use the ddp anymore. Thx~</p>

Liang: <p>For me, simply setting <code>find_unused_parameters</code> to <code>False</code> in DistributedDataParallel (DDP) solves the problem. There is no error anymore.</p>
<p>Ref: <a href="https://stackoverflow.com/questions/68000761/pytorch-ddp-finding-the-cause-of-expected-to-mark-a-variable-ready-only-once" rel="noopener nofollow ugc">https://stackoverflow.com/questions/68000761/pytorch-ddp-finding-the-cause-of-expected-to-mark-a-variable-ready-only-once</a></p>