Error on calling the backward pass of two losses

imhuim982: <p>I met a problem when I tried to call the backward function of two losses in sequential when using DDP.<br>
It works fine without DDP.<br>
I have a Function like this:</p>
<pre><code class="lang-auto">class WeightedSumFunc(torch.autograd.Function):
    a_ij_require_grad=True

    @staticmethod
    def forward(ctx, a_ij, e_ijk):
        ctx.save_for_backward(a_ij, e_ijk)
        output = torch.einsum('bi,bij-&gt;bj', a_ij, e_ijk)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # (N, dim)
        a_ij, e_ijk = ctx.saved_tensors
        if WeightedSumFunc.a_ij_require_grad:
            d_a_ij = torch.einsum('bj,bij-&gt;bi', grad_output, e_ijk)
        else:
            d_a_ij = None
        d_e_ijk = torch.einsum('bi,bj-&gt;bij', a_ij, grad_output)
        return d_a_ij, d_e_ijk
</code></pre>
<p>And I have two losses: loss1, and loss2. Here is my training code.</p>
<pre><code class="lang-auto">with torch.cuda.amp.autocast(dtype=amp_dtype):
    loss1, loss2 = self.model(**inputs)
    
    WeightedSumFunc.a_ij_require_grad = True
    scaler.scale(loss1 / accumulation_steps).backward(retain_graph=True)
    WeightedSumFunc.a_ij_require_grad = False
    scaler.scale(loss2 / accumulation_steps).backward()
</code></pre>
<p>I hope a_ij will backpropagate gradients during the backward pass in loss1 but not in loss2. Some parameters participate in both loss1 and loss2 and some don’t. I expect gradients can be accumulated after calling backward functions and it works fine without DDP.<br>
However, I receive a runtime error message when using DDP:<br>
“RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons: 1) Use of a module parameter outside the <code>forward</code> function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes. or try to use _set_static_graph() as a workaround if this module graph does not change during training loop.2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple <code>checkpoint</code> functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases in default. You can try to use _set_static_graph() as a workaround if your module graph does not change over iterations.”<br>
I didn’t use a parameter outside the forward function nor multiple ‘checkpoint’ functions.<br>
In my use-case, how can I accumulate gradients?<br>
Or is there any more elegant way to achieve the same goal?</p>
<p>Updates:<br>
I tried this</p>
<pre><code class="lang-auto">with model.no_sync():
      result = self.model(**inputs, coeff=coeff, temperature=temperature)
      if loss1 is not None:
          WeightedSumFunc.a_ij_require_grad = True
          scaler.scale(loss1 / accumulation_steps).backward(retain_graph=True)

      WeightedSumFunc.a_ij_require_grad = False
      scaler.scale(result.non_struct_loss / accumulation_steps).backward()
        
if (step + 1) % accumulation_steps == 0:
      for p in model.parameters():
          if p.grad is not None:
              torch.distributed.all_reduce(p.grad, op=torch.distributed.ReduceOp.SUM)
              p.grad /= torch.distributed.get_world_size()
</code></pre>
<p>Now no error is thrown, but I just find the loss is not as expected. The loss curve is still the same as training in a single card. How can I synchronize manually?</p>