offload fsdp state dict to cpu and rank 0 only

ad-freiburg · Jun 27, 2024 · 04f3be1 · 04f3be1
1 parent c633811
commit 04f3be1
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/python/text_utils/api/trainer.py b/python/text_utils/api/trainer.py
@@ -221,17 +221,16 @@ def __init__(
             # later, because FSDP handles mixed precision itself
             self.mixed_precision = None
 
-            offload_state_dict = self.info.world_size > 1
             FSDP.set_state_dict_type(
                 self.model,
                 StateDictType.FULL_STATE_DICT,
                 FullStateDictConfig(
-                    offload_to_cpu=offload_state_dict,
-                    rank0_only=offload_state_dict
+                    offload_to_cpu=True,
+                    rank0_only=True
                 ),
                 FullOptimStateDictConfig(
-                    offload_to_cpu=offload_state_dict,
-                    rank0_only=offload_state_dict
+                    offload_to_cpu=True,
+                    rank0_only=True
                 )
             )