Revert "Enable Unwrapping for Model State Dicts (FSDP) (#2959)" (#3096)

This reverts commit f35cbd1.
huggingface · Sep 10, 2024 · 1ce2eb6 · 1ce2eb6
1 parent 3fd02e6
commit 1ce2eb6
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -3334,8 +3334,6 @@ def get_state_dict(self, model, unwrap=True):
             from torch.distributed.fsdp import FullStateDictConfig, StateDictType
             from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
-            if unwrap:
-                model = self.unwrap_model(model)
             full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
             with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_state_dict_config):
                 state_dict = model.state_dict()