Support remote storage (fsspec URLs) for FSDP checkpoints

Author: zhixiangli

docs/source-fabric/guide/checkpoint/distributed_checkpoint.rst

@@ -48,6 +48,11 @@ The distributed checkpoint format is the default when you train with the :doc:`F
 With ``state_dict_type="sharded"``, each process/GPU will save its own file into a folder at the given path.
 This reduces memory peaks and speeds up the saving to disk.
 
+.. note::
+
+    The path can also be a remote filesystem URL supported by `fsspec <https://filesystem-spec.readthedocs.io/>`_, such as ``s3://my-bucket/checkpoint``, ``gs://my-bucket/checkpoint``, or ``abfs://my-container/checkpoint``.
+    This requires the corresponding fsspec implementation to be installed (e.g., ``s3fs``, ``gcsfs``, or ``adlfs``).
+
 .. collapse:: Full example
 
     .. code-block:: python
@@ -140,6 +145,8 @@ You can easily load a distributed checkpoint in Fabric if your script uses :doc:
 
 Note that you can load the distributed checkpoint even if the world size has changed, i.e., you are running on a different number of GPUs than when you saved the checkpoint.
 
+As with saving, the path can point to a remote filesystem URL supported by fsspec (e.g., ``s3://``, ``gs://``, ``abfs://``), and the checkpoint will be read directly from remote storage without needing to download it first.
+
 .. collapse:: Full example
 
     .. code-block:: python

docs/source-pytorch/advanced/model_parallel/fsdp.rst

@@ -400,6 +400,11 @@ The resulting checkpoint folder will have this structure:
 
 The “sharded” checkpoint format is the most efficient to save and load in Lightning.
 
+.. note::
+
+    The checkpoint path can also be a remote filesystem URL supported by `fsspec <https://filesystem-spec.readthedocs.io/>`_, such as ``s3://my-bucket/checkpoint``, ``gs://my-bucket/checkpoint``, or ``abfs://my-container/checkpoint``.
+    This requires the corresponding fsspec implementation to be installed (e.g., ``s3fs``, ``gcsfs``, or ``adlfs``).
+
 **Which checkpoint format should I use?**
 
 - ``state_dict_type="sharded"``: Use for pre-training very large models. It is fast and uses less memory, but it is less portable. An extra step is needed to :doc:`convert the sharded checkpoint into a regular checkpoint file <../../common/checkpointing_expert>`.
@@ -427,6 +432,8 @@ The Trainer will automatically recognize whether the provided path contains a ch
 Checkpoints saved with ``state_dict_type="full"`` can be loaded by all strategies, but sharded checkpoints can only be loaded by FSDP.
 Read :ref:`the checkpoints guide <checkpointing>` to explore more features.
 
+As with saving, the checkpoint path can point to a remote filesystem URL supported by fsspec (e.g., ``s3://``, ``gs://``, ``abfs://``), and the checkpoint will be read directly from remote storage without needing to download it first.
+
 
 ----
 

src/lightning/fabric/CHANGELOG.md

@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `using_sparse_model` and `sparse_cuda_acceleration_factor` parameters to `Throughput` so MFU defaults to the dense peak and opts into the sparse peak explicitly ([#21743](https://github.com/Lightning-AI/pytorch-lightning/pull/21743))
 
+- Added support for remote storage (fsspec URLs) when saving and loading distributed checkpoints with `FSDPStrategy` ([#21775](https://github.com/Lightning-AI/pytorch-lightning/pull/21775))
+
 ### Changed
 
 -

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,19 @@
     _Sharded,
     _validate_keys_for_strict_loading,
 )
+from lightning.fabric.utilities.cloud_io import (
+    _atomic_save,
+    _checkpoint_join,
+    _get_distributed_checkpoint_reader,
+    _get_distributed_checkpoint_writer,
+    _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 +451,8 @@ 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):
+        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 +473,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 +498,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 +517,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 +537,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 +578,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 +596,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 +608,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, weights_only=False)
 
             from lightning.fabric.strategies.model_parallel import (
                 _load_raw_module_state,
@@ -901,13 +909,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 +942,29 @@ 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 _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
         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
         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,18 @@ 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:
+        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

@@ -14,8 +14,10 @@
 """Utilities related to data saving/loading."""
 
 import errno
+import importlib
 import io
 import logging
+import shutil
 from pathlib import Path
 from typing import IO, Any, Optional, Union
 
@@ -176,3 +178,98 @@ 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)
+
+
+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)
+    FsspecWriter = _import_fsspec_dcp_filesystem("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)
+    FsspecReader = _import_fsspec_dcp_filesystem("FsspecReader")
+    return FsspecReader(path=str(path))
+
+
+def _import_fsspec_dcp_filesystem(name: str) -> Any:
+    """Import ``FsspecReader``/``FsspecWriter`` from torch's private DCP fsspec module.
+
+    These live in a private module that not every PyTorch build ships, so raise an actionable error
+    instead of letting a bare ``ImportError`` surface from deep in the call stack.
+
+    """
+    try:
+        module = importlib.import_module("torch.distributed.checkpoint._fsspec_filesystem")
+    except ImportError as e:
+        raise ImportError(
+            "Remote (fsspec) distributed checkpoints require"
+            " `torch.distributed.checkpoint._fsspec_filesystem`, which is not available in this"
+            " PyTorch build. Use a local checkpoint path or upgrade PyTorch."
+        ) from e
+    return getattr(module, name)

src/lightning/pytorch/CHANGELOG.md

@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `log_key_prefix` parameter to `LearningRateMonitor` callback for prefixing logged metric names ([#21612](https://github.com/Lightning-AI/pytorch-lightning/issues/21612))
 
+- Added support for remote storage (fsspec URLs) when saving and loading distributed checkpoints with `FSDPStrategy` ([#21775](https://github.com/Lightning-AI/pytorch-lightning/pull/21775))
+
 ### Changed
 
 -