diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 7e99c76ef..7f04c719a 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -977,7 +977,7 @@ def save_state(self, output_dir: str):
         output_dir = os.path.expanduser(output_dir)
         os.makedirs(output_dir, exist_ok=True)
         logger.info(f"Saving current state to {output_dir}")
-        weights = [self.get_state_dict(m) for m in self._models]
+        weights = [self.get_state_dict(m, unwrap=False) for m in self._models]
         save_location = save_accelerator_state(
             output_dir, weights, self._optimizers, self._schedulers, self.state.process_index, self.scaler
         )
@@ -1056,7 +1056,7 @@ def _get_devices(self, *args):
                         break
         return (model_device, optimizer_device)
 
-    def get_state_dict(self, model):
+    def get_state_dict(self, model, unwrap=True):
         is_zero_3 = False
         if self.distributed_type == DistributedType.DEEPSPEED:
             is_zero_3 = self.deepspeed_config["zero_optimization"]["stage"] == 3
@@ -1072,7 +1072,8 @@ def get_state_dict(self, model):
                     "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights."
                 )
         else:
-            model = self.unwrap_model(model)
+            if unwrap:
+                model = self.unwrap_model(model)
             state_dict = model.state_dict()
 
         if state_dict is not None:



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: Since this went unnoticed, might be good to add in a deepspeed test to the checkpointing test available here: https://github.com/huggingface/accelerate/blob/main/tests/test_state_checkpointing.py

If you think you can manage that feel free to add it to this PR, worst case I can add it in as a follow up
body: Thanks! I think I can add a test tomorrow.
body: Hi, although this solves the model weights issue, but now it fails when loading optimizer states.

Because deepspeed `optimizer.state_dict()` returns a dictionary.
https://github.com/microsoft/DeepSpeed/blob/9305916d6b858b20fca3ebe03940ceaaa13cdfbd/deepspeed/runtime/zero/stage_1_and_2.py#L1998

But deepspeed `optimizer.load_state_dict()`  expects a list of dictionary.
https://github.com/microsoft/DeepSpeed/blob/9305916d6b858b20fca3ebe03940ceaaa13cdfbd/deepspeed/runtime/zero/stage_1_and_2.py#L2135-L2167
because deepspeed `load_checkpoint` actually gets multiple state_dict (one for each partition or something) https://github.com/microsoft/DeepSpeed/blob/9305916d6b858b20fca3ebe03940ceaaa13cdfbd/deepspeed/runtime/engine.py#L2650

So I think I will just use this
```
if accelerator.distributed_type == DistributedType.DEEPSPEED:
    model.load_checkpoint(path)
    model.save_checkpoint(path)
else:
    accelerator.load_state(path)
    accelerator.save_state(path)
```
body: Since the test I tried to add fails when loading optimizer, I won't push it here. Please feel free to merge or close this PR.
body: Let's merge this first fix (thanks a lot!). Cc @pacman100 for the second test.