Fix resuming from checkpoint when using RayFSDPStrategy (#43594)

Signed-off-by: Paul Angerer <pangerer@canva.com>
ray-project · Mar 4, 2024 · b5aee36 · b5aee36
1 parent f79a9a7
commit b5aee36
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/python/ray/train/lightning/_lightning_utils.py b/python/ray/train/lightning/_lightning_utils.py
@@ -115,8 +115,16 @@ def lightning_module_state_dict(self) -> Dict[str, Any]:
                 ),
             ):
                 state_dict = self.model.state_dict()
+
+                ckpt_state_dict = {}
                 prefix_len = len("_forward_module.")
-                return {k[prefix_len:]: v for k, v in state_dict.items()}
+                for k, v in state_dict.items():
+                    if k.startswith("_forward_module."):
+                        non_prefixed_key = k[prefix_len:]
+                        ckpt_state_dict[non_prefixed_key] = v
+                    else:
+                        ckpt_state_dict[k] = v
+                return ckpt_state_dict
         else:
             # Otherwise Lightning uses Fairscale FSDP, no need to unshard by ourself.
             return super().lightning_module_state_dict()