accelerate cpu offload (#3242)

mosaicml · May 1, 2024 · b2edfbe · b2edfbe
1 parent 5eddaf3
commit b2edfbe
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/composer/core/state.py b/composer/core/state.py
@@ -890,7 +890,7 @@ def get_model_state_dict(self) -> Dict[str, Any]:
                 submodules=None,
                 options=StateDictOptions(
                     full_state_dict=self.fsdp_state_dict_type == 'full',
-                    cpu_offload=True,
+                    cpu_offload=self.fsdp_enabled,
                 ),
             )
         else:
@@ -930,7 +930,7 @@ def get_optim_state_dict(self) -> Dict[str, Any]:
                 submodules=None,
                 options=StateDictOptions(
                     full_state_dict=self.fsdp_state_dict_type == 'full',
-                    cpu_offload=True,
+                    cpu_offload=self.fsdp_enabled,
                 ),
             )
             return {type(optimizer).__qualname__: optim_state_dict}
@@ -1242,7 +1242,7 @@ def load_model_state(
                     options=StateDictOptions(
                         full_state_dict=self.fsdp_state_dict_type == 'full',
                         strict=strict,
-                        cpu_offload=True,
+                        cpu_offload=self.fsdp_enabled,
                     ),
                 )
             else:
@@ -1335,7 +1335,7 @@ def load_optim_state(self, state_dict: Dict[str, Any], strict: bool = True):
                     options=StateDictOptions(
                         full_state_dict=self.fsdp_state_dict_type == 'full',
                         strict=strict,
-                        cpu_offload=True,
+                        cpu_offload=self.fsdp_enabled,
                     ),
                 )
             else: