ITADN

Support remote/cloud storage (fsspec) for ModelParallelStrategy checkpointing

#21797Openzhixiangli 创建于 2026-07-02
featuredistributedcheckpointing
Z
zhixianglicommented
### 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 条评论