hij:

If i understand you correctly, registering it as a buffer only allows the stddev of rank 0 to be distributed to other processes.


Yes.

What I want to do is allow all tensors in all processes to contribute to calculating stddev.

You can do this by using the collective communication APIs (allreduce/allgather) in the forward pass. One caveat here is that, the collective communication API requires all processes in the same group to invoke the same API in the same order, otherwise, it might hang or the result might be wrong. If you are not sure whether the allreduce/allgather for stddev would interleave with other collectives, you can use the new_group 2 API to create a dedicated process group for collecting/summing stddev.