Support remote storage (fsspec URLs) for FSDP checkpoints

TAG=agy

CONV=e8cb1344-4be9-41e9-b673-24718e930cbc

Author: zhixiangli

src/lightning/fabric/strategies/fsdp.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import shutil
 import warnings
 from collections.abc import Generator
 from contextlib import AbstractContextManager, ExitStack, nullcontext
@@ -49,6 +48,17 @@
     _Sharded,
     _validate_keys_for_strict_loading,
 )
+from lightning.fabric.utilities.cloud_io import (
+    _atomic_save,
+    _checkpoint_join,
+    _is_checkpoint_dir,
+    _is_local_file_protocol,
+    _load,
+    _prepare_directory_checkpoint,
+    _remove_checkpoint,
+    _resolve_path,
+    get_filesystem,
+)
 from lightning.fabric.utilities.distributed import (
     ReduceOp,
     _distributed_is_initialized,
@@ -439,8 +449,9 @@ def save_checkpoint(
             )
 
         # broadcast the path from rank 0 to ensure all the states are saved in a common path
-        path = Path(self.broadcast(path))
-        if path.is_dir() and self._state_dict_type == "full" and not _is_sharded_checkpoint(path):
+        # remote URLs are kept as strings; only local paths are wrapped in `Path`
+        path = _resolve_path(self.broadcast(path))
+        if self._state_dict_type == "full" and _is_checkpoint_dir(path) and not _is_sharded_checkpoint(path):
             raise IsADirectoryError(f"The checkpoint path exists and is a directory: {path}")
 
         from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
@@ -461,9 +472,7 @@ def save_checkpoint(
         module = modules[0]
 
         if self._state_dict_type == "sharded":
-            if path.is_file():
-                path.unlink()
-            path.mkdir(parents=True, exist_ok=True)
+            _prepare_directory_checkpoint(path)
 
             state_dict_ctx = _get_sharded_state_dict_context(module)
 
@@ -488,11 +497,11 @@ def save_checkpoint(
             _distributed_checkpoint_save(converted_state, path)
 
             if self.global_rank == 0:
-                torch.save(metadata, path / _METADATA_FILENAME)
+                _atomic_save(metadata, _checkpoint_join(path, _METADATA_FILENAME))
 
         elif self._state_dict_type == "full":
             if _is_sharded_checkpoint(path):
-                shutil.rmtree(path)
+                _remove_checkpoint(path)
 
             state_dict_ctx = _get_full_state_dict_context(module, world_size=self.world_size)
             full_state: dict[str, Any] = {}
@@ -507,7 +516,7 @@ def save_checkpoint(
                     _apply_filter(key, filter or {}, converted, full_state)
 
             if self.global_rank == 0:
-                torch.save(full_state, path)
+                _atomic_save(full_state, path)
         else:
             raise ValueError(f"Unknown state_dict_type: {self._state_dict_type}")
 
@@ -527,7 +536,7 @@ def load_checkpoint(
                 " FSDPStrategy.load_checkpoint(..., state={'model': model, ...})"
             )
         # broadcast the path from rank 0 to ensure all the states are loaded from a common path
-        path = Path(self.broadcast(path))
+        path = _resolve_path(self.broadcast(path))
 
         if isinstance(state, Module):
             from lightning.fabric.strategies.model_parallel import _load_raw_module_state_from_path
@@ -568,11 +577,9 @@ def load_checkpoint(
                 module.load_state_dict(module_state[module_key], strict=strict)
 
                 if optimizers:
-                    from torch.distributed.checkpoint import FileSystemReader
-
                     # TODO: replace with newer APIs
                     # https://github.com/pytorch/pytorch/issues/119800#issuecomment-1942156271
-                    reader = FileSystemReader(path=path)
+                    reader = _get_distributed_checkpoint_reader(path)
                     # the optimizer states must be loaded separately
                     for optim_key, optim in optimizers.items():
                         optim_state = load_sharded_optimizer_state_dict(
@@ -588,7 +595,7 @@ def load_checkpoint(
                         optim.load_state_dict(flattened_osd)
 
             # Load metadata (anything not a module or optimizer)
-            metadata = torch.load(path / _METADATA_FILENAME, weights_only=weights_only)
+            metadata = _load(_checkpoint_join(path, _METADATA_FILENAME), weights_only=weights_only)
             requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
             _validate_keys_for_strict_loading(requested_metadata_keys, metadata.keys(), strict=strict)
             for key in requested_metadata_keys:
@@ -600,7 +607,7 @@ def load_checkpoint(
             return metadata
 
         if _is_full_checkpoint(path):
-            checkpoint = _lazy_load(path)
+            checkpoint = _lazy_load(path) if _is_local_file_protocol(str(path)) else _load(path)
 
             from lightning.fabric.strategies.model_parallel import (
                 _load_raw_module_state,
@@ -901,13 +908,19 @@ def _get_full_state_dict_context(
     return state_dict_type_context  # type: ignore[return-value]
 
 
-def _is_sharded_checkpoint(path: Path) -> bool:
+def _is_sharded_checkpoint(path: _PATH) -> bool:
     """A heuristic check to determine whether the path points to a directory with checkpoint shards."""
-    return path.is_dir() and (path / _METADATA_FILENAME).is_file()
+    if _is_local_file_protocol(str(path)):
+        path = Path(path)
+        return path.is_dir() and (path / _METADATA_FILENAME).is_file()
+    fs = get_filesystem(path)
+    return fs.isfile(str(path).rstrip("/") + "/" + _METADATA_FILENAME)
 
 
-def _is_full_checkpoint(path: Path) -> bool:
-    return path.is_file()
+def _is_full_checkpoint(path: _PATH) -> bool:
+    if _is_local_file_protocol(str(path)):
+        return Path(path).is_file()
+    return get_filesystem(path).isfile(str(path))
 
 
 def _has_fsdp_modules(module: object) -> TypeGuard[Module]:
@@ -928,38 +941,54 @@ def _move_torchmetrics_to_device(module: torch.nn.Module, device: torch.device)
         metric.to(device)  # `.to()` is in-place
 
 
-def _distributed_checkpoint_save(converted_state: dict[str, Any], path: Path) -> None:
+def _get_distributed_checkpoint_writer(path: _PATH) -> Any:
+    if _is_local_file_protocol(str(path)):
+        from torch.distributed.checkpoint import FileSystemWriter
+
+        # FSDP's FileSystemWriter streams the tensors to disk to minimize memory peaks
+        return FileSystemWriter(path=path, single_file_per_rank=True)
+    from torch.distributed.checkpoint._fsspec_filesystem import FsspecWriter
+
+    return FsspecWriter(path=str(path), single_file_per_rank=True)
+
+
+def _get_distributed_checkpoint_reader(path: _PATH) -> Any:
+    if _is_local_file_protocol(str(path)):
+        from torch.distributed.checkpoint import FileSystemReader
+
+        return FileSystemReader(path=path)
+    from torch.distributed.checkpoint._fsspec_filesystem import FsspecReader
+
+    return FsspecReader(path=str(path))
+
+
+def _distributed_checkpoint_save(converted_state: dict[str, Any], path: _PATH) -> None:
     if _TORCH_GREATER_EQUAL_2_3:
         from torch.distributed.checkpoint import save
 
-        # let torch automatically infer the writer to use. This might also support fsspec paths in the future
-        # https://github.com/pytorch/pytorch/issues/118036
+        # torch automatically infers the writer to use from the checkpoint_id, including
+        # FsspecWriter for remote (fsspec) URLs. https://github.com/pytorch/pytorch/issues/118036
         save(converted_state, checkpoint_id=path)
     else:  # deprecated
-        from torch.distributed.checkpoint import FileSystemWriter
-
         if _TORCH_GREATER_EQUAL_2_2:
             from torch.distributed.checkpoint import save
         else:
             from torch.distributed.checkpoint import save_state_dict as save
-        # FSDP's FileSystemWriter streams the tensors to disk to minimize memory peaks
-        writer = FileSystemWriter(path=path, single_file_per_rank=True)
+        writer = _get_distributed_checkpoint_writer(path)
         save(converted_state, writer)
 
 
-def _distributed_checkpoint_load(module_state: dict[str, Any], path: Path) -> None:
+def _distributed_checkpoint_load(module_state: dict[str, Any], path: _PATH) -> None:
     if _TORCH_GREATER_EQUAL_2_3:
         from torch.distributed.checkpoint import load
 
-        # let torch automatically infer the reader to use. This might also support fsspec paths in the future
-        # https://github.com/pytorch/pytorch/issues/118036
+        # torch automatically infers the reader to use from the checkpoint_id, including
+        # FsspecReader for remote (fsspec) URLs. https://github.com/pytorch/pytorch/issues/118036
         load(module_state, checkpoint_id=path)
     else:  # deprecated
-        from torch.distributed.checkpoint import FileSystemReader
-
         if _TORCH_GREATER_EQUAL_2_2:
             from torch.distributed.checkpoint import load
         else:
             from torch.distributed.checkpoint import load_state_dict as load
-        reader = FileSystemReader(path=path)
+        reader = _get_distributed_checkpoint_reader(path)
         load(module_state, reader)

src/lightning/fabric/strategies/model_parallel.py

@@ -43,6 +43,7 @@
     _BackwardSyncControl,
     _validate_keys_for_strict_loading,
 )
+from lightning.fabric.utilities.cloud_io import _is_local_file_protocol, _load
 from lightning.fabric.utilities.distributed import (
     ReduceOp,
     _distributed_is_initialized,
@@ -529,15 +530,19 @@ def _has_dtensor_modules(module: object) -> TypeGuard[Module]:
     return isinstance(module, Module) and any(isinstance(t, DTensor) for t in module.parameters())
 
 
-def _load_raw_module_state_from_path(path: Path, module: Module, world_size: int, strict: bool = True) -> None:
+def _load_raw_module_state_from_path(path: _PATH, module: Module, world_size: int, strict: bool = True) -> None:
     """Loads the state dict from a file path into the FSDP module."""
     if not _is_full_checkpoint(path):
         raise ValueError(
             "Failed to load checkpoint directly into the model. The given path must be a single file containing the"
             f" full state dict: {path}"
         )
-    # Use `lazy_load`/`mmap` instead to avoid storing a copy of the full checkpoint per rank
-    state_dict = torch.load(path, mmap=True, map_location="cpu") if _TORCH_GREATER_EQUAL_2_3 else _lazy_load(path)
+    if _is_local_file_protocol(str(path)):
+        # Use `lazy_load`/`mmap` instead to avoid storing a copy of the full checkpoint per rank
+        state_dict = torch.load(path, mmap=True, map_location="cpu") if _TORCH_GREATER_EQUAL_2_3 else _lazy_load(path)
+    else:
+        # mmap cannot be used over object storage; read the object through fsspec
+        state_dict = _load(path, map_location="cpu")
     _load_raw_module_state(state_dict=state_dict, module=module, world_size=world_size, strict=strict)
 
 

src/lightning/fabric/utilities/cloud_io.py

@@ -16,6 +16,7 @@
 import errno
 import io
 import logging
+import shutil
 from pathlib import Path
 from typing import IO, Any, Optional, Union
 
@@ -176,3 +177,61 @@ def _is_dir(fs: AbstractFileSystem, path: Union[str, Path], strict: bool = False
 
 def _is_local_file_protocol(path: _PATH) -> bool:
     return fsspec.utils.get_protocol(str(path)) == "file"
+
+
+def _resolve_path(path: _PATH) -> Union[str, Path]:
+    """Return a ``Path`` for local file paths and a plain ``str`` for remote fsspec URLs.
+
+    ``Path()`` collapses the double slash in a URL (e.g. ``gs://bucket`` -> ``gs:/bucket``),
+    corrupting it, so remote URLs must be kept as strings.
+
+    """
+    if _is_local_file_protocol(str(path)):
+        return Path(path)
+    return str(path)
+
+
+def _checkpoint_join(path: Union[str, Path], name: str) -> Union[str, Path]:
+    """Join ``name`` onto a checkpoint ``path`` without corrupting remote URLs."""
+    if isinstance(path, Path):
+        return path / name
+    return str(path).rstrip("/") + "/" + name
+
+
+def _is_checkpoint_dir(path: Union[str, Path]) -> bool:
+    """Return whether ``path`` points to an existing directory, supporting fsspec paths."""
+    if isinstance(path, Path):
+        return path.is_dir()
+    return get_filesystem(path).isdir(str(path))
+
+
+def _prepare_directory_checkpoint(path: Union[str, Path]) -> None:
+    """Ensure ``path`` is a directory for a sharded checkpoint.
+
+    Removes a conflicting file sitting at ``path`` and creates the directory. Creating a
+    directory is a no-op on object storage, which has no real directories.
+
+    """
+    if isinstance(path, Path):
+        if path.is_file():
+            path.unlink()
+        path.mkdir(parents=True, exist_ok=True)
+        return
+    fs = get_filesystem(path)
+    if fs.isfile(str(path)):
+        fs.rm(str(path))
+    if not _is_object_storage(fs):
+        fs.makedirs(str(path), exist_ok=True)
+
+
+def _remove_checkpoint(path: Union[str, Path]) -> None:
+    """Remove a checkpoint file or directory (recursively), supporting fsspec paths."""
+    if isinstance(path, Path):
+        if path.is_dir():
+            shutil.rmtree(path)
+        elif path.exists():
+            path.unlink()
+        return
+    fs = get_filesystem(path)
+    if fs.exists(str(path)):
+        fs.rm(str(path), recursive=True)