FSDP Strategy checkpoint loading

[xin-w8023]

Description & Motivation

pytorch-lightning/src/lightning/pytorch/strategies/fsdp.py

Lines 587 to 596 in b9680a3

	def load_checkpoint(self, checkpoint_path: _PATH) -> Dict[str, Any]:
	# broadcast the path from rank 0 to ensure all the states are loaded from a common path
	path = Path(self.broadcast(checkpoint_path))
	from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
	assert self.model is not None
	assert self.lightning_module is not None
	if _is_sharded_checkpoint(path):

Here we can only check a local path whether sharded or not. Any remote file path like hdfs://, the function will eventually raise an error.

Pitch

No response

Alternatives

No response

Additional context

No response

cc @Borda