Support remote/cloud storage (fsspec) for ModelParallelStrategy checkpointing
featuredistributedcheckpointing
### Description & Motivation
#21779 identified that FSDP checkpointing assumed a local filesystem and fixed it by routing saves/loads through cloud-aware helpers in cloud_io.py (_resolve_path, _checkpoint_join, _is_checkpoint_dir, _prepare_directory_checkpoint, _remove_checkpoint).
ModelParallelStrategy (Fabric and PyTorch) has the identical problem, since it shares the same DCP-based save/load path as FSDP:
1. Path corruption — Path(self.broadcast(path)) collapses gs://bucket/... to gs:/bucket/..., producing an invalid remote URL.
2. Local-only file ops — path.is_dir(), path.mkdir(), path.unlink(), and shutil.rmtree() are used directly and don't work against object storage.
3. mmap on remote storage — full-checkpoint loading unconditionally passes mmap=True to torch.load, which requires local file access and fails over object storage.
### Pitch
Apply the same fix from #21779 to ModelParallelStrategy:
- In both lightning/fabric/strategies/model_parallel.py and lightning/pytorch/strategies/model_parallel.py, replace Path(...)/shutil/pathlib calls with the cloud_io helpers (_resolve_path, _checkpoint_join, _is_checkpoint_dir, _prepare_directory_checkpoint, _remove_checkpoint) already introduced for FSDP.
- Guard the full-checkpoint load path so mmap=True is only used for local files (via _is_local_file_protocol); fall back to _load(path, map_location="cpu", weights_only=...) for remote URLs.
- Add regression coverage asserting a gs://...-style URL reaches the DCP save/load calls uncorrupted (not truncated to gs:/).
### Alternatives
_No response_
### Additional context
_No response_
cc @lantiga @justusschock
0 条评论