register hook for gradients in convnext dtensor training example to get better performance

KingsleyLiu-NV · KingsleyLiu-NV · commit b5927a518baf · 2023-11-19T21:33:26.000-08:00
diff --git a/torch/distributed/_tensor/examples/convnext_example.py b/torch/distributed/_tensor/examples/convnext_example.py
@@ -156,6 +156,9 @@ def _conv_fn(
     for name, param in module.named_parameters():
         dist_spec = [Replicate()]
         dist_param = torch.nn.Parameter(distribute_tensor(param, device_mesh, dist_spec))
+        dist_param.register_hook(
+            lambda grad: grad.redistribute(placements=dist_spec)
+        )
         name = "_".join(name.split("."))
         module.register_parameter(name, dist_param)
 
