update

Author: awaelchli

src/lightning/pytorch/strategies/fsdp.py

@@ -53,6 +53,7 @@
     _TORCH_GREATER_EQUAL_2_0,
 )
 from lightning.fabric.utilities.init import _EmptyInit
+from lightning.fabric.utilities.load import _lazy_load
 from lightning.fabric.utilities.optimizer import _optimizers_to_device
 from lightning.fabric.utilities.seed import reset_seed
 from lightning.fabric.utilities.types import _PATH, ProcessGroup, ReduceOp
@@ -572,9 +573,8 @@ def load_checkpoint(self, checkpoint_path: _PATH) -> Dict[str, Any]:
             return metadata
 
         if _is_full_checkpoint(path):
-            # TODO: Support lazy-loading here (see Fabric)
-            checkpoint = torch.load(path, map_location="cpu")
-            _load_raw_module_state(checkpoint["state_dict"], world_size=self.world_size, module=self.model)
+            checkpoint = _lazy_load(path) if _TORCH_GREATER_EQUAL_2_0 else torch.load(path, map_location="cpu")
+            _load_raw_module_state(checkpoint["state_dict"], module=self.model, world_size=self.world_size)
 
             from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
             from torch.distributed.fsdp import OptimStateKeyType