Enable Unwrapping for Model State Dicts (FSDP) (huggingface#2959)

Signed-off-by: Alex-Brooks <[email protected]>
khanhi2r · Jul 31, 2024 · f35cbd1 · f35cbd1
1 parent a14260c
commit f35cbd1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -3289,6 +3289,8 @@ 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()