Manually edit tensor grad before DDP gradient sync

albert.cwkuo: <p>I wonder if it is valid to manually edit a tensor’s grad of a DDP model before syncing the gradient. This is what I am trying to do:</p>
<pre><code class="lang-auto">1. base_model = MyModel().cuda(device_ids[0])
2. ddp_model = DDP(base_model, device_ids)

3. outputs = ddp_model(input)
4. loss1, loss2 = loss_fn(outputs)

5. with ddp_model.no_sync():
6.     local_loss = loss1 + 0. * loss2
7.     local_loss.backward(retain_graph=True)
8.     for p in base_model.sub_model2.parameters():
9.         p.grad *= 0.
10.     for p in base_model.sub_model3.parameters():
11.         p.grad *= 0.

12. local_loss = 0. * loss1 + loss2
13. local_loss.backward()

14. optimizer.step()
</code></pre>
<p>As shown in the code snippet above, I manually modify the base_model’s gradient at lines 8 - 11 before syncing the gradient at line 13. My goal is to use loss1 to only update sub_model1 in the base_model, and use loss2 to update the whole base_model.</p>
<p>The code runs without error, but I am concerned if this manual modification of tensor gradient will cause any issue to the gradient sync mechanism in DDP.</p>

mrshenli: <p>Hey <a class="mention" href="/u/albert.cwkuo">@albert.cwkuo</a></p>
<p>With the above code, I think DDP still syncs all grads for both <code>loss1</code> and <code>loss2</code>, because the flag controlled by <code>no_sync</code> ctx manager is used when calling <code>DistributedDataParallel.forward()</code>. So, as the forward is out of the <code>no_sync</code> context, DDP would still prepare to sync all grads during the backward pass.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/5036c94a6e868963e0354fc04c92e204d8d77677/torch/nn/parallel/distributed.py#L477-L498" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/5036c94a6e868963e0354fc04c92e204d8d77677/torch/nn/parallel/distributed.py#L477-L498" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/5036c94a6e868963e0354fc04c92e204d8d77677/torch/nn/parallel/distributed.py#L477-L498</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="477" style="counter-reset: li-counter 476 ;">
<li>@contextmanager</li>
<li>def no_sync(self):</li>
<li>    r"""</li>
<li>    A context manager to disable gradient synchronizations across DDP</li>
<li>    processes. Within this context, gradients will be accumulated on module</li>
<li>    variables, which will later be synchronized in the first</li>
<li>    forward-backward pass exiting the context.</li>
<li>
</li><li>    Example::</li>
<li>
</li><li>        &gt;&gt;&gt; ddp = torch.nn.DistributedDataParallel(model, pg)</li>
<li>        &gt;&gt;&gt; with ddp.no_sync():</li>
<li>        ...   for input in inputs:</li>
<li>        ...     ddp(input).backward()  # no synchronization, accumulate grads</li>
<li>        ... ddp(another_input).backward()  # synchronize grads</li>
<li>    """</li>
<li>    old_require_backward_grad_sync = self.require_backward_grad_sync</li>
<li>    self.require_backward_grad_sync = False</li>
<li>    try:</li>
<li>        yield</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/5036c94a6e868963e0354fc04c92e204d8d77677/torch/nn/parallel/distributed.py#L477-L498" target="_blank" rel="nofollow noopener">show original</a>

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

<p>How is <code>MyModel</code> implemented. Does it contain two independent submodules <code>sub_model1</code> and <code>sub_model2</code> and do sth like the following?</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        self.sub_model1 = SomeModel1()
        self.sub_model2 = SomeModel2()

    def forward(self, input):
        return self.sub_model1(input), self.sub_model2(input)
</code></pre>

albert.cwkuo: <p>Thanks <a class="mention" href="/u/mrshenli">@mrshenli</a> for your reply. This is what <code>MyModel</code> do internally.</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        self.sub_model1 = SomeModel1()
        self.sub_model2 = SomeModel2()

    def forward(self, input):
        out1 = self.sub_model2(input)
        out2 = self.sub_model1(out1)
        return out1, out2
</code></pre>
<p>I want the gradient to be synced eventually when I call <code>backward()</code> at line 11, so the above code seems correct? My goal is to accumulate gradient from both loss1 and loss2 to sub_model2 and accumulate gradient only from loss2 to sub_model1. That’s why I try to zero out grad in sub_model1.</p>
<p>Note that I use <code>local_loss = loss1 + 0.0 * loss2</code> and <code>local_loss = 0.0 * loss1 + loss2</code> to mask out part of the loss before calling <code>local_loss.backward()</code>.</p>
<p>----------------update------------</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        self.sub_model1 = SomeModel1()
        self.sub_model2 = SomeModel2()
        self.sub_model3 = SomeModel3()

    def forward(self, input):
        out1 = self.sub_model2(self.sub_model1(input))
        out2 = self.sub_model3(out1)
        return out1, out2
</code></pre>
<p>My goal is to accumulate gradient from both loss1 and loss2 to sub_model1 and accumulate gradient only from loss2 to sub_model2 and sub_model3. That’s why I try to zero out grad in sub_model2 and sub_model3.</p>

mrshenli: <aside class="quote no-group" data-username="albert.cwkuo" data-post="3" data-topic="87042">
<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/albert.cwkuo/48/12567_2.png" class="avatar"> albert.cwkuo:</div>
<blockquote>
<p>My goal is to accumulate gradient from both loss1 and loss2 to sub_model2 and accumulate gradient only from loss2 to sub_model1.</p>
</blockquote>
</aside>
<p>In that case, calling backward once on <code>loss1+loss2</code> might be sufficient? Is the following result what you want?</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

class MyModel(nn.Module):

    def __init__(self):
        super().__init__()
        with torch.no_grad():
            self.net1 = nn.Linear(1, 1)
            self.net1.weight.copy_(torch.ones(1, 1))
            self.net1.bias.copy_(torch.zeros(1))
            self.net2 = nn.Linear(1, 1)
            self.net2.weight.copy_(torch.ones(1, 1))
            self.net2.bias.copy_(torch.zeros(1))

    def forward(self, x):
        out1 = self.net1(x)
        out2 = self.net2(out1)
        return out1 + out2


print("==============")
model = MyModel()
model(torch.ones(1, 1)).sum().backward()
print("net1 grad is: ", model.net1.weight.grad)
print("net2 grad is: ", model.net2.weight.grad)

print("==============")
model = MyModel()
model.net1(torch.ones(1, 1)).sum().backward()
print("net1 grad is: ", model.net1.weight.grad)
print("net1 grad is: ", model.net2.weight.grad)

print("==============")
model = MyModel()
model.net2(torch.ones(1, 1)).sum().backward()
print("net1 grad is: ", model.net1.weight.grad)
print("net2 grad is: ", model.net2.weight.grad)
</code></pre>
<p>outputs are</p>
<pre><code class="lang-auto">==============
net1 grad is:  tensor([[2.]])
net2 grad is:  tensor([[1.]])
==============
net1 grad is:  tensor([[1.]])
net2 grad is:  None
==============
net1 grad is:  None
net2 grad is:  tensor([[1.]])
</code></pre>

albert.cwkuo: <p>Okay that makes sense, but let me clarify how MyModel works internally so that the proposed solution may not work in the case. The reply has been updated.</p>
<p>In short, I have 3 subnets inside MyModel. One of the loss depends on sub_model1 and sub_model2, but I only want to update sub_model1 with this loss. Therefore I need to zero out the grad in sub_model2 when calling backward on that loss.</p>

mrshenli: <aside class="quote no-group" data-username="albert.cwkuo" data-post="3" data-topic="87042">
<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/albert.cwkuo/48/12567_2.png" class="avatar"> albert.cwkuo:</div>
<blockquote>
<p>My goal is to accumulate gradient from both loss1 and loss2 to sub_model1 and accumulate gradient only from loss2 to sub_model2 and sub_model3.</p>
</blockquote>
</aside>
<p>With the above statement, should the forward function be something like below?</p>
<pre data-code-wrap="python"><code class="lang-python">    def forward(self, input):
        out1 = self.sub_model1(input)
        out2 = self.sub_model3(self.sub_model2(out1))
        return out1, out2
</code></pre>
<p>With this code, <code>out1.sum().backward()</code> will only compute grads for <code>sub_model1</code>, and <code>out2.sum().backward()</code> will compute grads for all sub-models. And <code>(out1 + out2).sum().backward()</code> should meet the cited statement above.</p>

albert.cwkuo: <p>In my use case it’s</p>
<pre><code class="lang-python">    def forward(self, input):
        out1 = self.sub_model2(self.sub_model1(input))
        out2 = self.sub_model3(out1)
        return out1, out2
</code></pre>
<p>That’s the tricky part. That’s why I want to zero out the gradient of sub_mode2 after calling backward on loss1 (loss1 is computed on out1).</p>