Glad it works!



 11123:

So based on my experiment, I think maybe your work can also solve the distributed gpu grad_fn gathering problem? like in Will “dist.all_gather” break the auto gradient graph?.


Yes it seems that diffdist can handle that case. Of course different processes will have different computational graphs but with diffdist some nodes are inserted in the graph that will cause them to sync and communicate with each other. For example, doing a Send operation will cause a Recv to be called during backward in order to receive the gradient.