I am trying to understand the use case here:

I have read that all_gather do not retain the gradient information. For my application, I want stddev to be part of the computation graph. How would I go about doing this?

IIUC, stddev in this case is an intermediate output in forward and it’s not a model parameter? So you need its gradient during the backward pass to compute parameter gradients, but you don’t need to retain its gradient for optimizer step()? If above is true, why existing parameter gradient synchronization in DDP not sufficient?

And could you also point me to an example/tutorial for these usage.

Sure, below is the tutorial, and please search for “All-Reduce example.”
https://pytorch.org/tutorials/intermediate/dist_tuto.html 6