hij:

However, since I have limited gpu memory, I can only train with batch size of 1. And there isn’t a point to calculate stddev over a batchsize of 1 even if I do DDP. So in this case, would the allgather allreduce work?


I see. I am not sure if the result would still be correct in this case even if allgather and allreduce can retain gradients. IIUC, if this is trained without DDP (assume there are large enough GPU memory), then both feats and stddev are calculated based on all inputs. When trained with DDP, feats are now only derived from local inputs, and you would like to have stddev to be based on global inputs. So, when you cat feats and stddev, the output of the forward now represents a different thing. I am not sure if the loss function can handle that. Even if it can, what does averaging gradient mean in this case?
If above (local feats + global stddev) is the expected behavior, there might be a few ways to implement this.


Implement a custom autograd function. E.g., its forward function can use an allgather to collect stddev from all processes, and its backward function can use another allgather to collect gradients and then extract the part belongs to the local stddev and sum them up.


Use torch.distributed.rpc 2. There can be a master and a few workers, where each worker calculates the feats and stddev for its own input, and then the master gathers all feats and stddev to compute the final loss. Some more tutorials for this:
a. https://pytorch.org/tutorials/intermediate/rpc_tutorial.html 2
b. https://pytorch.org/tutorials/intermediate/rpc_param_server_tutorial.html 1