diff --git a/src/fairseq2/nn/data_parallel/_fsdp1.py b/src/fairseq2/nn/data_parallel/_fsdp1.py index ed48750da..114bbf403 100644 --- a/src/fairseq2/nn/data_parallel/_fsdp1.py +++ b/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 diff --git a/src/fairseq2/nn/data_parallel/_fsdp2.py b/src/fairseq2/nn/data_parallel/_fsdp2.py index a9cdba1ec..203aeecea 100644 --- a/src/fairseq2/nn/data_parallel/_fsdp2.py +++ b/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