Support remote storage (fsspec URLs) for FSDP checkpoints

[zhixiangli]

What does this PR do?

Fix #21779

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

[zhixiangli]

Hi @tchaton @justusschock @ethanwharris

I'm eager to get this PR in shape, but I'm currently having trouble accessing the CI failure logs.

When I follow the check links, I hit a login page. I attempted to register for a Lightning AI account to view them, but the phone verification step is consistently rejecting my Singapore phone number with this error:

Invalid number (support code: 09332104)

I have already reached out to support@lightning.ai with my phone details. Could someone kindly let me know if registering for an account is the correct path for external contributors to view these logs, or if special access needs to be granted in another way?

If anyone on the team is able to look into the support ticket or safely expedite the verification process, I would immensely appreciate it.

Thank you so much for your time and help!

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 96.80851% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 87%. Comparing base (4819088) to head (6b2de49).
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #21775   +/-   ##
=======================================
  Coverage      87%      87%           
=======================================
  Files         270      270           
  Lines       24010    24064   +54     
=======================================
+ Hits        20787    20874   +87     
+ Misses       3223     3190   -33     

[Copilot]

Pull request overview

This PR extends Lightning’s FSDP checkpointing (Fabric + PyTorch) to support remote storage paths expressed as fsspec URLs (e.g., s3://, gs://, memory://) by avoiding pathlib.Path URL corruption and routing filesystem operations through fsspec-aware helpers.

Changes:

• Added fsspec-aware checkpoint path utilities in cloud_io.py (safe join, dir checks, prepare/remove checkpoint, DCP reader/writer selection).
• Updated Fabric and PyTorch FSDP strategies (and model-parallel loading) to use these utilities for remote checkpoint save/load, including DCP metadata handling.
• Added regression tests + documentation + changelog entries for remote checkpoint paths.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/tests_pytorch/strategies/test_fsdp.py	Adds regression coverage ensuring remote URLs aren’t corrupted during PyTorch FSDP sharded checkpoint save.
tests/tests_fabric/utilities/test_cloud_io.py	Adds unit tests for new fsspec-aware path and checkpoint filesystem helpers.
tests/tests_fabric/strategies/test_fsdp.py	Adds Fabric FSDP tests for remote paths, DCP reader/writer selection, and remote full-checkpoint loading behavior.
src/lightning/pytorch/strategies/fsdp.py	Uses fsspec-aware path handling and filesystem ops for distributed checkpointing in PyTorch FSDP strategy.
src/lightning/pytorch/CHANGELOG.md	Notes new remote-storage support for FSDP distributed checkpoints.
src/lightning/fabric/utilities/cloud_io.py	Introduces new helper functions for resolving/joining/removing/checkpointing paths + DCP fsspec reader/writer import.
src/lightning/fabric/strategies/model_parallel.py	Adjusts model-parallel load path to avoid mmap/lazy loading for remote checkpoints.
src/lightning/fabric/strategies/fsdp.py	Implements remote-path support across save/load for Fabric FSDP distributed checkpoints.
src/lightning/fabric/CHANGELOG.md	Notes new remote-storage support for FSDP distributed checkpoints.
docs/source-pytorch/advanced/model_parallel/fsdp.rst	Documents that FSDP checkpoint paths can be fsspec URLs and lists required fsspec backends.
docs/source-fabric/guide/checkpoint/distributed_checkpoint.rst	Documents that distributed checkpoint paths can be fsspec URLs and can be loaded directly from remote storage.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[deependujha]

We need to update load_sharded_optimizer_state_dict for torch>=2.3.

/home/zeus/miniconda3/envs/cloudspace/lib/python3.12/site-packages/torch/distributed/checkpoint/optimizer.py:348: FutureWarning: `load_state_dict` is deprecated and will be removed in future versions. Please use `load` instead.

I'll handle it in a separate PR. verified with PL on T4 4gpu with saving sharded checkpoints on AWS S3. ✅

Code used

from typing import override
import torch
import warnings
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
import lightning as L
from lightning.pytorch.strategies import FSDPStrategy

L.seed_everything(42)

warnings.filterwarnings(
    "ignore",
    message=".*Please use DTensor instead.*",
    category=FutureWarning,
)
warnings.filterwarnings(
    "ignore",
    message=".*_get_pg_default_device.*",
)

class MySuperCoolDataModule(L.LightningDataModule):
    def __init__(
        self,
        batch_size=32,
        dataset_size=64,
        input_dim=8,
    ):
        super().__init__()
        self.save_hyperparameters()

        self.batch_size = batch_size
        self.x = torch.randn(dataset_size, input_dim)
        self.y = torch.randn(dataset_size)

    def setup(self, stage=None):
        self.dataset = TensorDataset(self.x, self.y)

    def train_dataloader(self):
        return DataLoader(
            self.dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=2,
        )


class MySuperCoolModel(L.LightningModule):
    def __init__(self, input_dim=8):
        super().__init__()
        self.save_hyperparameters()

        self.layers = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 2048),
            nn.ReLU(),
            nn.Linear(2048, 2048),
            nn.ReLU(),
            nn.Linear(2048, 2048),
            nn.ReLU(),
            nn.Linear(2048, 256),
            nn.ReLU(),
            nn.Linear(256, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
        )

    def forward(self, x):
        return self.layers(x)

    def training_step(self, batch, batch_idx):
        x, y = batch

        y_hat = self(x).squeeze(-1)
        loss = nn.functional.mse_loss(y_hat, y)

        self.log("train_loss", loss)

        return loss

    @override
    def configure_optimizers(self):
        return torch.optim.Adam(
            self.parameters(),
            lr=1e-3,
        )


def main():
    model = MySuperCoolModel()

    datamodule = MySuperCoolDataModule(
        batch_size=32,
        dataset_size=64,
        input_dim=8,
    )

    policy = {nn.Linear,}

    trainer = L.Trainer(
        default_root_dir="s3://lightning-checkpoints/simple-checkpoint-02/",
        accelerator="auto",
        devices="auto",
        strategy=FSDPStrategy(auto_wrap_policy=policy, state_dict_type="sharded"),
        precision="32-true",
        max_epochs=2,
        log_every_n_steps=1,
    )

    trainer.fit(model, datamodule)

    trainer.print("-" * 80)
    trainer.print("Training done. ✅")
    # trainer.print(torch.cuda.memory_summary())


if __name__ == "__main__":
    main()