Unified storage_options support for FSDPStrategy, ModelParallelStrategy, and TorchCheckpointIO

Description & Motivation

When saving and loading distributed checkpoints (e.g., via PyTorch Distributed Checkpoint / DCP with FSDPStrategy or ModelParallelStrategy), PyTorch allows configuring storage backend parameters such as:

• thread_count: Number of concurrent I/O threads per rank to save shards (both torch.distributed.checkpoint.FileSystemWriter and torch.distributed.checkpoint._fsspec_filesystem.FsspecWriter default to thread_count=1).
• per_thread_copy_ahead: Buffer size for pre-staging data before writing.
• fsspec kwargs: Authentication tokens, endpoints, and caching parameters when saving directly to remote cloud buckets (e.g., S3, GCS).

Currently, PyTorch Lightning and Fabric hardcode default instantiation without passing storage kwargs. Furthermore, passing storage_options to FSDPStrategy.save_checkpoint, ModelParallelStrategy.save_checkpoint, or TorchCheckpointIO.save_checkpoint raises a TypeError.

Pitch

Enable storage_options: Optional[dict[str, Any]] = None universally across strategies and I/O utilities:

1. Strategy Initialization: Allow configuring default storage_options at strategy initialization (e.g. FSDPStrategy(storage_options={"thread_count": 8}) or ModelParallelStrategy(storage_options={"thread_count": 8})).
2. Strategy & Fabric Save/Load: Allow per-call overrides via strategy.save_checkpoint(..., storage_options=...), Fabric.save(..., storage_options=...), and Trainer.save_checkpoint(..., storage_options=...).
3. Internal Forwarding:
   • Forward storage_options to _get_distributed_checkpoint_writer(path, **storage_options) (which passes thread_count, per_thread_copy_ahead, etc. to FileSystemWriter / FsspecWriter).
   • Forward storage_options to _get_distributed_checkpoint_reader(path, **storage_options) (FileSystemReader / FsspecReader).
   • Forward storage_options to _atomic_save(..., storage_options=...) and _load(..., storage_options=...) for fsspec cloud saving of full checkpoints and meta.pt.

User Experience / Code Example

from lightning.pytorch import Trainer
from lightning.pytorch.strategies import FSDPStrategy

# 1. Configure thread count and I/O settings at strategy level
strategy = FSDPStrategy(
    state_dict_type="sharded",
    storage_options={"thread_count": 8, "per_thread_copy_ahead": 20_000_000},
)
trainer = Trainer(strategy=strategy, ...)

# 2. Or override at save/load time
trainer.save_checkpoint("checkpoints/epoch=10", storage_options={"thread_count": 16})

# 3. With Fabric
from lightning.fabric import Fabric

fabric = Fabric(strategy="fsdp")
fabric.save("checkpoints/step=1000", state, storage_options={"thread_count": 8})

Alternatives

No response

Additional context

Historical Context & Why TypeError Was Originally Added

In PR #11891 (commit d31126c33), storage_options was introduced to Trainer.save_checkpoint and CheckpointIO.save_checkpoint to allow custom CheckpointIO plugins to receive arbitrary storage parameters. To prevent silent discarding of arguments when users passed storage_options to default plugins (TorchCheckpointIO, FSDPStrategy, ModelParallelStrategy), an explicit TypeError was raised:

if storage_options is not None:
    raise TypeError("... is not supported because ... does not use CheckpointIO.")

Today, Lightning's cloud_io and distributed strategies natively support fsspec filesystems and PyTorch DCP storage backends (FileSystemWriter, FsspecWriter, FileSystemReader, FsspecReader). Raising TypeError prevents users from configuring essential performance parameters like thread count and filesystem credentials.

cc @lantiga