Save replicated items on all shards (#1247)

Author: zyaoj

src/fairseq2/nn/data_parallel/_fsdp1.py

@@ -13,7 +13,6 @@
 from typing import TypeAlias
 
 import torch
-import torch.distributed as dist
 from torch import Tensor
 from torch.distributed import ProcessGroup
 from torch.distributed._shard.sharded_tensor import ShardedTensor
@@ -209,22 +208,17 @@ def fsdp1_local_state_dict(module: FSDP1) -> dict[str, object]:
             action="ignore", message=r".*Please use DTensor instead.*"
         )
 
-        sdp_rank = dist.get_rank(module.process_group)
-
         for name, item in module.state_dict().items():
             if isinstance(item, ShardedTensor):
                 local_shards = item.local_shards()
                 if not local_shards:
                     continue  # means the tensor is sharded unevenly.
 
-                state_dict[name] = item.local_tensor().detach()
-            # Save replicated items only on the first intra-node (i.e. sharded)
-            # gang.
-            elif sdp_rank == 0:
-                if isinstance(item, Tensor):
-                    item = item.detach()
+                item = item.local_tensor().detach()
+            elif isinstance(item, Tensor):
+                item = item.detach()
 
-                state_dict[name] = item
+            state_dict[name] = item
 
     return state_dict
 

src/fairseq2/nn/data_parallel/_fsdp2.py

@@ -195,36 +195,15 @@ def wrap(module: Module, reshard_after_forward: bool | None = None) -> FSDP2:
 def fsdp2_local_state_dict(module: FSDP2) -> dict[str, object]:
     sharded_state_dict = module.state_dict()
 
-    device_mesh = None
-
-    for value in sharded_state_dict.values():
-        if isinstance(value, DTensor):
-            device_mesh = value.device_mesh
-
-            break
-
-    if device_mesh is not None:
-        try:
-            sdp_rank = device_mesh.get_local_rank(mesh_dim="intra")
-        except KeyError:
-            raise ValueError(
-                "The device mesh of `module` does not have a dimension named 'intra'."
-            ) from None
-    else:
-        sdp_rank = 0
-
     state_dict: dict[str, object] = {}
 
     for key, value in sharded_state_dict.items():
         if isinstance(value, DTensor):
-            state_dict[key] = cast(DTensor, value.detach()).to_local()
-        # Save replicated items only on the first intra-node (i.e. sharded)
-        # gang.
-        elif sdp_rank == 0:
-            if isinstance(value, Tensor):
-                value = value.detach()
-
-            state_dict[key] = value
+            value = cast(DTensor, value.detach()).to_local()
+        elif isinstance(value, Tensor):
+            value = value.detach()
+
+        state_dict[key] = value
 
     return state_dict