Repository: pytorch/pytorch

Issue:
Allow Full State Dict with 2D FSDP + TP
### 🐛 Describe the bug

Torch does not allow 2D FSDP + TP to get FULL_STATE_DICT. However, if I remove checks here:
https://github.com/pytorch/pytorch/blob/3f62b05d31d4b29d60874b05adc0e5aedbad3722/torch/distributed/fsdp/_state_dict_utils.py#L303-L310
and warnings here:
https://github.com/pytorch/pytorch/blob/3f62b05d31d4b29d60874b05adc0e5aedbad3722/torch/distributed/fsdp/_state_dict_utils.py#L811-L815
it seems to work fine.

Is there a reason we need these checks? Can we remove them given things seem to work fine?

### Versions

Nightly.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225


Discussion:
@pytorchbot label "oncall: distributed"
Hey @mvpatel2000. Thanks for raising the issue. Initially we wanted to disable it because the full_state_dict is only full state dict on data parallel dimension. So to avoid the discrepancy in naming and semantics, we disable it to let users just use SHARDED_STATE_DICT in 2D case. 

Just to make sure that we understand your use case, so you do want to save the full FSDP state_dict (but still sharded on TP dimension), right? 

cc. @fegin 
@wz337 hm... in our case it seems to gather everything (including TP dimension)... Why would it not do that?
> @wz337 hm... in our case it seems to gather everything (including TP dimension)... Why would it not do that?

Hmm. Interesting... We didn't enable the full_state_dict path and nor did we test it before. iirc, the full_state_dict path currently is not aware of 2D and DTensorExtension. 

Hmm. You are saying without making any change, simply removing the assert gives you the full_state_dict as you expect? 
I think dcp's get_state_dict need to be able to get a full_state_dict if user want to? the full_state_dict option on the code pointer is from FSDP, so it might be FSDP only without aware of TP?
> Hmm. You are saying without making any change, simply removing the assert gives you the full_state_dict as you expect?
Yep!

We have some model which is FSDP + TPed. We then want to gather the model in order to package it up for easy sharing, e.g. into a Huggingface checkpoint. You can see how we do this here:
https://github.com/mosaicml/llm-foundry/blob/main/llmfoundry/callbacks/hf_checkpointer.py#L230

Currently, we do:
```
use_full_state_dict = not state.is_model_ddp and isinstance(state_dict_model, FSDP)
state_dict_context = fsdp_state_dict_type_context(
  original_model, state_dict_type='full') if use_full_state_dict else contextlib.nullcontext()

with state_dict_context:
  state_dict = state_dict_model.state_dict()
```

With the new API, we can do:
```
from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict

state_dict = get_model_state_dict(
  state_dict_model,
  options=StateDictOptions(
    full_state_dict=True,
    cpu_offload=True,
  ),
)
```
and this seems to work fine! We then create a new model which does not have TP applied, load in the model state dict, and export via Huggingface APIs
Oh ok. 

> > Hmm. You are saying without making any change, simply removing the assert gives you the full_state_dict as you expect?
> > Yep!
> 
> We have some model which is FSDP + TPed. We then want to gather the model in order to package it up for easy sharing, e.g. into a Huggingface checkpoint. You can see how we do this here: https://github.com/mosaicml/llm-foundry/blob/main/llmfoundry/callbacks/hf_checkpointer.py#L230
> 
> Currently, we do:
> 
> ```
> use_full_state_dict = not state.is_model_ddp and isinstance(state_dict_model, FSDP)
> state_dict_context = fsdp_state_dict_type_context(
>   original_model, state_dict_type='full') if use_full_state_dict else contextlib.nullcontext()
> 
> with state_dict_context:
>   state_dict = state_dict_model.state_dict()
> ```
> 
> With the new API, we can do:
> 
> ```
> from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
> 
> state_dict = get_model_state_dict(
>   state_dict_model,
>   options=StateDictOptions(
>     full_state_dict=True,
>     cpu_offload=True,
>   ),
> )
> ```
> 
> and this seems to work fine! We then create a new model which does not have TP applied, load in the model state dict, and export via Huggingface APIs

Oh. Ok. Thanks for the code pointer. So to clarify, there are two layers here, FSDP state dict and DCP's distributed state dict (I am referring to `get_model_state_dict`). 

When we call the generic `get_model_state_dict` with full_state_dict, it calls into `_gather_state_dict`, which will all_gather the full tensor for you regardless of whether it is 1D or 2D. This is the reason that you are seeing the correct result. https://github.com/pytorch/pytorch/blob/e2ee87d48bd2a70a18920c800da9fd0522be619d/torch/distributed/_state_dict_utils.py#L171

If you use the context manager full_state_dict for FSDP and remove the `RuntimeError`, it would give you 1D DTensor state_dict. So in order to avoid confusion, we disable it in FSDP's full_state_dict originally.

As long as users are using `get_model_state_dict` to do state_dict retrieval, I think I am fine with removing the warning and RuntimeError. @fegin Wdyt? 

Ah got it, thanks for the clarification! 

It'd be great if we could remove the error/warning then as suggested. It makes our life a LOT easier to easily export models. 
> Ah got it, thanks for the clarification!
> 
> It'd be great if we could remove the error/warning then as suggested. It makes our life a LOT easier to easily export models.

Kk. Thanks for reporting the issue. I don't see any issue of removing the RuntimeError and Warning. Just created the PR to let CI run and will get the team to review it and provide feedback if any. 
@wz337  there is one more issue though. [here](https://github.com/pytorch/pytorch/blob/4cf6d1172b25a0a7b773b40cd2401b9a6dad0fc1/torch/distributed/checkpoint/state_dict.py#L399-L401) it only offloads to CPU **AFTER** getting **ALL** the state_dict. It easily gets OOM if there are huge DTensor in the state_dict. Can we run the following code each time we get the DTensor in state_dict ? I reported this issue [here](https://github.com/pytorch/pytorch/issues/121062). 
```
return _gather_state_dict(
      state_dict, cpu_offload=info.cpu_offload, ranks_only=ranks_only
)
```

We had to add the following hook to avoid OOM before calling `get_model_state_dict `

```
cpu_offload = True
def dtensor_to_tensor_hook(
    module: nn.Module,
    state_dict: Dict[str, Any],
    prefix: str,
    *args: Any,
) -> Dict[str, Any]:
    dtensor_fqns = []
    for fqn in state_dict.keys():
        tensor = state_dict[fqn]
        if isinstance(tensor, DTensor):
            dtensor_fqns.append(fqn)
            tensor = tensor.full_tensor()
            if dist.get_global_rank() == 0:
                if cpu_offload:
                    tensor = tensor.cpu()
                state_dict[fqn] = tensor
    if dist.get_global_rank() != 0:
        for fqn in dtensor_fqns:
            del state_dict[fqn]
    return state_dict

hooks = []
for name, module in state_dict_model.named_modules():
    if isinstance(module, FSDP):
        hooks.append(
            module._register_state_dict_hook(
                dtensor_to_tensor_hook))

```

Resolving patch:
diff --git a/torch/distributed/fsdp/_optim_utils.py b/torch/distributed/fsdp/_optim_utils.py
--- a/torch/distributed/fsdp/_optim_utils.py
+++ b/torch/distributed/fsdp/_optim_utils.py
@@ -2083,10 +2083,5 @@ def _set_optim_use_dtensor(
                 "DeviceMesh is not compatible with LOCAL_STATE_DICT.",
                 "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict.",
             )
-        elif state_dict_type == StateDictType.FULL_STATE_DICT:
-            logger.warning(
-                "Found both state_dict_type FULL_STATE_DICT and device_mesh. "  # noqa: G004
-                "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict."
-            )
         else:
             state_dict_settings.optim_state_dict_config._use_dtensor = True
diff --git a/torch/distributed/fsdp/_state_dict_utils.py b/torch/distributed/fsdp/_state_dict_utils.py
--- a/torch/distributed/fsdp/_state_dict_utils.py
+++ b/torch/distributed/fsdp/_state_dict_utils.py
@@ -292,11 +292,6 @@ def _full_pre_state_dict_hook(
     """
     if getattr(fsdp_state, "_device_mesh", False):
         parent_mesh = _mesh_resources.get_parent_mesh(fsdp_state._device_mesh)
-        if parent_mesh:
-            raise RuntimeError(
-                f"Found FSDP's device_mesh {fsdp_state._device_mesh} has a parent device_mesh {parent_mesh}.",
-                "We do not support FULL_STATE_DICT for 2D FSDP + TP. Please use FSDP SHARDED_STATE_DICT instead.",
-            )
 
     _common_pre_state_dict_hook(module, fsdp_state)
     _common_unshard_pre_state_dict_hook(
@@ -804,11 +799,6 @@ def _set_use_dtensor(fsdp_state: _FSDPState) -> None:
                 "DeviceMesh is not compatible with LOCAL_STATE_DICT.",
                 "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict.",
             )
-        elif state_dict_type == StateDictType.FULL_STATE_DICT:
-            logger.warning(
-                "Found both state_dict_type FULL_STATE_DICT and device_mesh. "  # noqa: G004
-                "Please set state_dict_type to SHARDED_STATE_DICT to get DTensor state_dict."
-            )
         else:
             fsdp_state._state_dict_config._use_dtensor = True
 


Test patch:
diff --git a/test/distributed/fsdp/test_fsdp_dtensor_state_dict.py b/test/distributed/fsdp/test_fsdp_dtensor_state_dict.py
--- a/test/distributed/fsdp/test_fsdp_dtensor_state_dict.py
+++ b/test/distributed/fsdp/test_fsdp_dtensor_state_dict.py
@@ -313,30 +313,6 @@ def test_raises_warning_or_errors(self):
             with FSDP.state_dict_type(model, StateDictType.LOCAL_STATE_DICT):
                 optim_state_dict = FSDP.optim_state_dict(model, optim)
 
-        with self.assertLogs(
-            "torch.distributed.fsdp._state_dict_utils", level="WARNING"
-        ) as log:
-            with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
-                state_dict = model.state_dict()
-                self.assertEqual(len(log.records), 1)
-                self.assertEqual(len(log.output), 1)
-                self.assertIn(
-                    "Found both state_dict_type FULL_STATE_DICT and device_mesh.",
-                    log.output[0],
-                )
-
-        with self.assertLogs(
-            "torch.distributed.fsdp._optim_utils", level="WARNING"
-        ) as log:
-            with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
-                state_dict = FSDP.optim_state_dict(model, optim)
-                self.assertEqual(len(log.records), 1)
-                self.assertEqual(len(log.output), 1)
-                self.assertIn(
-                    "Found both state_dict_type FULL_STATE_DICT and device_mesh.",
-                    log.output[0],
-                )
-
 
 instantiate_parametrized_tests(TestFSDPWithDeviceMeshAndDTensor)
 if __name__ == "__main__":