I did some test. Tensor.register_hook() does work in DDP. I first wrap my model with DDP, then I modify the grad manually with a custom function in the forward code where I construct my model:

image761×24 6.1 KB
Here’s the gradients after calling loss.backward():
If don’t modify the grad:

And if modify the grad:

You can see the gradients are just 3 times after being modified, which means .register_hook() works well in DDP.