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:
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()

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.
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.