Repository: huggingface/accelerate
Pull request: fix fsdp state_dict_config because of PyTorch changes
### What does this PR do?
1. Fixes #831

Problem statement:
FSDP tests fail due to ImportError _state_dict_type_to_config  (has been moved inside of set_state_dict_type method
### System Info

```Shell
- `Accelerate` version: 0.15.0.dev0
- Platform: Linux-5.15.68.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
- Python version: 3.9.13
- Numpy version: 1.23.4
- PyTorch version (GPU?): 1.14.0.dev20221107 (True)
- `Accelerate` default config:
        Not found
```


### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [ ] My own task or dataset (give details below)

### Reproduction

`pytest tests/fsdp/test_fsdp.py`

results in

```
FAILED tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_auto_wrap_policy - ImportError: cannot import name '_state_dict_type_to_config' from 'torch.distributed.fsdp.fully_sharded_data_parall...
FAILED tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_backward_prefetch - ImportError: cannot import name '_state_dict_type_to_config' from 'torch.distributed.fsdp.fully_sharded_data_parall...
FAILED tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_cpu_offload - ImportError: cannot import name '_state_dict_type_to_config' from 'torch.distributed.fsdp.fully_sharded_data_parall...
FAILED tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_mixed_precision - ImportError: cannot import name '_state_dict_type_to_config' from 'torch.distributed.fsdp.fully_sharded_data_parall...
FAILED tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_sharding_strategy - ImportError: cannot import name '_state_dict_type_to_config' from 'torch.distributed.fsdp.fully_sharded_data_parall...
FAILED tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_state_dict_type - ImportError: cannot import name '_state_dict_type_to_config' from 'torch.distributed.fsdp.fully_sharded_data_parall...
```

### Expected behavior

```Shell
The tests to pass.

Note that the fail is due to a change in recent pytorch, the variable _state_dict_type_to_config has been moved into the static method set_state_dict_type , so isn't available for direct import from fully_sharded_data_parallel.py
```

Hints from issue discussion:
cc @pacman100 
Hello @Thomas-MMJ, Thank you for bringing this to the notice. Please use PyTorch 1.13.0 stable version for time-being. We will resolve this issue once PyTorch version 1.14.0 has a stable release.

Commit message: fix fsdp test

Patch to review:
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index bb14078d4e2..57ead9be633 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -31,6 +31,7 @@
 import torch
 
 from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH, FSDP_STATE_DICT_TYPE, MODEL_NAME, OPTIMIZER_NAME
+from .versions import is_torch_version
 
 
 class KwargsHandler:
@@ -526,9 +527,9 @@ def __post_init__(self):
         from torch.distributed.fsdp.fully_sharded_data_parallel import (
             BackwardPrefetch,
             CPUOffload,
+            FullStateDictConfig,
             ShardingStrategy,
             StateDictType,
-            _state_dict_type_to_config,
         )
 
         if self.sharding_strategy is None:
@@ -549,12 +550,8 @@ def __post_init__(self):
             state_dict_type_policy = os.environ.get("FSDP_STATE_DICT_TYPE", "FULL_STATE_DICT")
             self.state_dict_type = StateDictType(FSDP_STATE_DICT_TYPE.index(state_dict_type_policy) + 1)
 
-            if self.state_dict_type == StateDictType.FULL_STATE_DICT:
-                self.state_dict_config = _state_dict_type_to_config[self.state_dict_type](
-                    offload_to_cpu=True, rank0_only=True
-                )
-            else:
-                self.state_dict_config = _state_dict_type_to_config[self.state_dict_type]()
+            if self.state_dict_type == StateDictType.FULL_STATE_DICT and self.state_dict_config is None:
+                self.state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
 
     @staticmethod
     def get_module_class_from_name(module, name):
@@ -616,9 +613,14 @@ def save_model(self, accelerator, model, output_dir, model_index=0):
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
         from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 
-        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+        if is_torch_version("<=", "1.13.0"):
             with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
                 state_dict = model.state_dict()
+        else:
+            FSDP.set_state_dict_type(model, self.state_dict_type, self.state_dict_config)
+            state_dict = model.state_dict()
+
+        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
             weights_name = f"{MODEL_NAME}.bin" if model_index == 0 else f"{MODEL_NAME}_{model_index}.bin"
             output_model_file = os.path.join(output_dir, weights_name)
             if accelerator.process_index == 0:
@@ -626,8 +628,6 @@ def save_model(self, accelerator, model, output_dir, model_index=0):
                 torch.save(state_dict, output_model_file)
                 print(f"Model saved to {output_model_file}")
         else:
-            with FSDP.state_dict_type(model, self.state_dict_type, self.state_dict_config):
-                state_dict = model.state_dict()
             weights_name = (
                 f"{MODEL_NAME}_rank{accelerator.process_index}.bin"
                 if model_index == 0
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
index 249d2b69277..e90106c17bb 100644
--- a/tests/fsdp/test_fsdp.py
+++ b/tests/fsdp/test_fsdp.py
@@ -95,7 +95,7 @@ def test_backward_prefetch(self):
                     self.assertEqual(fsdp_plugin.backward_prefetch, BackwardPrefetch(i + 1))
 
     def test_state_dict_type(self):
-        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType, _state_dict_type_to_config
+        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 
         for i, state_dict_type in enumerate(FSDP_STATE_DICT_TYPE):
             env = self.dist_env.copy()
@@ -103,9 +103,6 @@ def test_state_dict_type(self):
             with mockenv_context(**env):
                 fsdp_plugin = FullyShardedDataParallelPlugin()
                 self.assertEqual(fsdp_plugin.state_dict_type, StateDictType(i + 1))
-                self.assertEqual(
-                    type(fsdp_plugin.state_dict_config), type(_state_dict_type_to_config[StateDictType(i + 1)]())
-                )
                 if state_dict_type == "FULL_STATE_DICT":
                     self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu)
                     self.assertTrue(fsdp_plugin.state_dict_config.rank0_only)


Code review comments:
diff_hunk: @@ -616,18 +613,21 @@ def save_model(self, accelerator, model, output_dir, model_index=0):
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
         from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
 
-        if self.state_dict_type == StateDictType.FULL_STATE_DICT:
+        if is_torch_version("<=", "1.13.0"):
path: src/accelerate/utils/dataclasses.py
text: @user1:
```suggestion
        if is_torch_version("<=", "1.13.5"):
```
Just in case they make a patch release before 1.14 is out.