Enable checkpoint with FSDP on single device (Lightning-AI#124)

sateeshs · Apr 12, 2023 · 7b2f995 · 7b2f995
1 parent 5d98989
commit 7b2f995
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lit_llama/utils.py b/lit_llama/utils.py
@@ -15,7 +15,7 @@ def save_model_checkpoint(fabric, model, file_path):
     """
 
     if isinstance(fabric.strategy, FSDPStrategy):
-        save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        save_policy = FullStateDictConfig(offload_to_cpu=(fabric.world_size > 1), rank0_only=True)
         with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
             state_dict = model._forward_module.state_dict()
     else: