feat: enable storage_options support in FSDP and ModelParallel strategies

[Yonghui-Lee]

What does this PR do?

Fixes #21905

This PR enables unified support for the storage_options parameter across FSDPStrategy, ModelParallelStrategy, TorchCheckpointIO, and Fabric.save/Fabric.load. It removes the legacy TypeError guards and wires storage_options directly through to the underlying DCP storage writers/readers, fsspec, and atomic saving utilities.

Key Changes

1. lightning.fabric.utilities.cloud_io:

   • Updated _get_distributed_checkpoint_writer and _get_distributed_checkpoint_reader to accept **kwargs and pass them to FileSystemWriter/FsspecWriter and FileSystemReader/FsspecReader (enabling thread_count, per_thread_copy_ahead, etc.).
   • Updated _load and _atomic_save to accept storage_options: Optional[dict[str, Any]] = None and forward to get_filesystem and fsspec.core.url_to_fs.

2. lightning.fabric.plugins.io.torch_io:

   • Removed the TypeError guard in TorchCheckpointIO.save_checkpoint and forwarded storage_options to get_filesystem and _atomic_save.

3. FSDPStrategy (Fabric & PyTorch):

   • Added storage_options: Optional[dict[str, Any]] = None to FSDPStrategy.__init__.
   • Forwarded storage_options in save_checkpoint and load_checkpoint to _distributed_checkpoint_save, _distributed_checkpoint_load, _get_distributed_checkpoint_reader, and _atomic_save.

4. ModelParallelStrategy (Fabric & PyTorch):

   • Added storage_options: Optional[dict[str, Any]] = None to ModelParallelStrategy.__init__.
   • Forwarded storage_options in save_checkpoint, load_checkpoint, _save_checkpoint, and _load_checkpoint.

5. Fabric & Utilities:

   • Added storage_options parameter to Fabric.save and Fabric.load.
   • Added storage_options parameter to _load_distributed_checkpoint.

6. Tests:

   • Updated existing tests that previously asserted TypeError on storage_options across Fabric and PyTorch test suites (test_fsdp.py, test_model_parallel.py, test_trainer_checkpoint.py).
   • Added assertions ensuring storage_options (e.g. thread_count) are passed through to _distributed_checkpoint_save, _atomic_save, FileSystemWriter, and FsspecWriter.

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