How does DistributedDataParallel handle parameters whose requires_grad flag is False?

username: ayalaa2
cooked: <p>Hello, I’m trying to figure out what happen behind the scenes in DistributedDataParallel when it comes to parameters that do not require a gradient. I cannot find a clear answer on this in the documentations.</p>
<p>Assume we have three layers: A --&gt; B --&gt; C. Suppose that A and B both have their <code>requires_grad</code> flag as False. If this model is wrapped in DistributedDataParallel, <strong>will there be any process communication that needs to be done in the backward pass of layer A or C?</strong> Specifically with the sharing of gradients.</p>
<p>My problem is that I have a large model and I am extremely bottle necked by process communication. I would like to freeze some of my layers such that less gradients need to be shared. I understand that depending on my model, the computation cost may be the same but I really need to bring down the communication cost.</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/ayalaa2">@ayalaa2</a>, DistributedDataParallel’s (DDP) ctor would go through all parameters and skip the ones whose <code>requires_grad=False</code>. So, there won’t be communication on those grad, but you will have to set their <code>require_grad</code> field before passing it to DDP. After the ctor, changing the <code>requires_grad</code> attribute makes no difference. See the code below in DDP ctor.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/6bd88f581a6323d026e08b65ffee75bfe162501f/torch/nn/parallel/distributed.py#L456-L464" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/6bd88f581a6323d026e08b65ffee75bfe162501f/torch/nn/parallel/distributed.py#L456-L464" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/6bd88f581a6323d026e08b65ffee75bfe162501f/torch/nn/parallel/distributed.py#L456-L464</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="456" style="counter-reset: li-counter 455 ;">
<li># Build tuple of (module, parameter) for all parameters that require grads.</li>
<li>modules_and_parameters = [</li>
<li>    [</li>
<li>        (module, parameter)</li>
<li>        for module in replica.modules()</li>
<li>        for parameter in filter(</li>
<li>            lambda parameter: parameter.requires_grad,</li>
<li>            parameters(module, recurse=False))</li>
<li>    ] for replica in self._module_copies]</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Another way to skip communication is to use the <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel.no_sync" rel="nofollow noopener"><code>no_sync</code></a> context manager.</p>

username: wydwww
cooked: <p>Hi <a class="mention" href="/u/mrshenli">@mrshenli</a>, you mentioned that DDP can skip the gradient communication for parameters whose <code>requires_grad=False</code> but the flag must be set before wrapping the model with DDP.</p>
<p>I have 2 following questions:</p>
<ol>
<li>
<p>If I set <code>requires_grad=False</code> during training after the DDP ctor, will these parameters be updated anymore, if they still conduct communication?</p>
</li>
<li>
<p>Is there a way to <strong>dynamically</strong> freeze more parameters and skip their communication <strong>during DDP training</strong> after constructing the DDP model?</p>
</li>
</ol>
<p>My use case is gradually freezing more layers during training/fine-tuning. Thanks!</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/wydwww">@wydwww</a></p>
<aside class="quote no-group" data-username="wydwww" data-post="3" data-topic="90736">
<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/wydwww/48/12800_2.png" class="avatar"> wydwww:</div>
<blockquote>
<p>If I set <code>requires_grad=False</code> during training after the DDP ctor, will these parameters be updated anymore, if they still conduct communication?</p>
</blockquote>
</aside>
<p>DDP only builds the communication buffer once (it’s actually twice: 1. using the reverse order of <code>model.parameters</code>, 2. using the autograd order of the first iteration). So DDP will still communicate all parameter’s gradients, even if some of them are marked as <code>requires_grad=False</code> after DDP ctor.</p>
<blockquote>
<ol>
<li>Is there a way to <strong>dynamically</strong> freeze more parameters and skip their communication <strong>during DDP training</strong> after constructing the DDP model?</li>
</ol>
</blockquote>
<p>If this is not very frequent (say, once per epoch), you can destroy the DDP instance and create a new one.</p>

username: wydwww
cooked: <p>Thanks. I will try this solution.</p>

username: palo_bajo
cooked: <p>I found this suggestion quite useful. I am facing exactly the same problem described here but I can’t fix it by just destroying the DDP and creating a new instance <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_frowning_face.png?v=12" title=":slightly_frowning_face:" class="emoji" alt=":slightly_frowning_face:" loading="lazy" width="20" height="20"> Basically I have a model wrapped around DDP, and after 10 epochs I want to freeze some parts using <code>requires_grad = False</code> and keep training the rest of the model.</p>
<p>To destroy the DDP model instance I am following <a href="https://discuss.pytorch.org/t/safely-removing-a-module-from-ddp/45753" class="inline-onebox">Safely removing a Module from DDP</a>, this is, <code>model = model.module</code>, but still not working.</p>
<p>Any suggestions?</p>

username: kumpera
cooked: <p>Could you go in more detal of how is it exactly failing?<br>
Do you have logs of the failure to help troubleshoot it?</p>

username: wydwww
cooked: <p>Did you first extract the module from DDP (<code>model = model.module</code>), then create a new DDP with the extracted <code>model</code>? And what do you mean by not working?</p>