[FSDP2] Follow-up fix to correct relaxed overlap test (pytorch#132953)

The previous PR forgot to include dummy all-gathers before backward, so the reference time was too short, causing the test to still fail.

I verified the test passes locally.

This should close pytorch#120961 (again).

Pull Request resolved: pytorch#132953
Approved by: https://github.com/weifengpy
ghstack dependencies: pytorch#132869

test/distributed/_composable/fsdp/test_fully_shard_overlap.py

@@ -108,6 +108,14 @@ def ref_fwd_bwd():
                     ]
                     dist.all_gather_into_tensor(dummy_ag_output, dummy_ag_input)
                 loss = ref_model(inp).sum()
+                # Run dummy all-gathers per weight again since we are
+                # resharding after forward
+                for lin in ref_model:
+                    dummy_ag_output = torch.empty_like(lin.weight)
+                    dummy_ag_input = torch.chunk(dummy_ag_output, self.world_size)[
+                        self.rank
+                    ]
+                    dist.all_gather_into_tensor(dummy_ag_output, dummy_ag_input)
                 loss.backward()
                 # Run dummy reduce-scatters per weight
                 for lin in ref_model: