In this case, I’m assuming that A is a part of the network M, although during backwards, since according to DDP, A went unused, it is marked as ready for reduction, but then we attempt to re-mark A’s parameters when it gets grads in the backwards pass.
Is it possible to create a wrapper network that wraps steps 2 and 3 into a single module and returns both loss1 and loss2/the values needed to compute it? That seems like it would avoid the double mark issue.
Also, the version of PyTorch you are using and a script to repro the issue would also be valuable.