Update on "[FSDP2] Computed grad divide factors at runtime"

Andrew Gu · Andrew Gu · commit 05f56970a750 · 2024-05-03T11:19:10.000-07:00
cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k

[ghstack-poisoned]
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_init.py b/test/distributed/_composable/fsdp/test_fully_shard_init.py
@@ -738,51 +738,54 @@ def test_hsdp_broadcast_across_replicas(self):
         )
         model_args = ModelArgs()
         model = Transformer(model_args)
+        # Add a buffer to show that this flow works for buffers too
+        model.register_buffer("buf", torch.randn((model_args.dim,)))
         for module in model.modules():
             if isinstance(module, TransformerBlock):
                 fully_shard(module, mesh=mesh)
         fully_shard(model, mesh=mesh)
 
-        # Only preserve the model parameters on the replicate mesh's rank 0
+        # Only preserve the model states on the replicate mesh's rank 0
         if mesh.get_local_rank("replicate") > 0:
-            print(f"[Rank {self.rank}] filling with 1337")
-            for param in model.parameters():
-                param.detach().fill_(1337)
+            for tensor in itertools.chain(model.parameters(), model.buffers()):
+                tensor.detach().fill_(1337)
 
         # Check that replicas are different
-        for param_name, param in model.named_parameters():
-            local_param = param.to_local()
-            local_param_list = [
-                torch.empty_like(local_param) for _ in range(mesh["replicate"].size())
+        for tensor in itertools.chain(model.parameters(), model.buffers()):
+            local_tensor = tensor.to_local() if isinstance(tensor, DTensor) else tensor
+            local_tensor_list = [
+                torch.empty_like(local_tensor) for _ in range(mesh["replicate"].size())
             ]
             dist.all_gather(
-                local_param_list, local_param, group=mesh.get_group("replicate")
+                local_tensor_list, local_tensor, group=mesh.get_group("replicate")
             )
-            for other_local_param in local_param_list[1:]:
-                self.assertEqual(other_local_param.shape, local_param_list[0].shape)
-                self.assertNotEqual(other_local_param, local_param_list[0])
+            for other_local_tensor in local_tensor_list[1:]:
+                self.assertEqual(other_local_tensor.shape, local_tensor_list[0].shape)
+                self.assertNotEqual(other_local_tensor, local_tensor_list[0])
 
         # Broadcast from replicate mesh's rank 0
         replicate_group = mesh.get_group("replicate")
-        for param in model.parameters():
+        for tensor in itertools.chain(model.parameters(), model.buffers()):
             # E.g. for mesh [[0, 1, 2, 3], [4, 5, 6, 7]] sharding on dim-1 and
             # replicating on dim-0, broadcast with sources 0, 1, 2, 3
             src_rank = dist.get_process_group_ranks(replicate_group)[0]
             torch.distributed.broadcast(
-                param.to_local(), src=src_rank, group=replicate_group
+                tensor.to_local() if isinstance(tensor, DTensor) else tensor,
+                src=src_rank,
+                group=replicate_group,
             )
 
         # Check that replicas are the same
-        for param_name, param in model.named_parameters():
-            local_param = param.to_local()
-            local_param_list = [
-                torch.empty_like(local_param) for _ in range(mesh["replicate"].size())
+        for tensor in itertools.chain(model.parameters(), model.buffers()):
+            local_tensor = tensor.to_local() if isinstance(tensor, DTensor) else tensor
+            local_tensor_list = [
+                torch.empty_like(local_tensor) for _ in range(mesh["replicate"].size())
             ]
             dist.all_gather(
-                local_param_list, local_param, group=mesh.get_group("replicate")
+                local_tensor_list, local_tensor, group=mesh.get_group("replicate")
             )
-            for other_local_param in local_param_list[1:]:
-                self.assertEqual(other_local_param, local_param_list[0])
+            for other_local_tensor in local_tensor_list[1:]:
+                self.assertEqual(other_local_tensor, local_tensor_list[0])
 
         # Check that we can run an iteration without erroring
         inp = torch.randint(0, model_args.vocab_size, (2, 16), device="cuda")
