Lightning-AI / Lightning-AI/pytorch-lightning
fabric FSDP strategy save/load checkpoint does not support s3 url
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
Other ParallelStrategy implementations support saving and loading checkpoints from an S3 URL using fsspec. Fabric’s FSDP strategy also uses torch.load and torch.save, which support S3 URLs as well. However, the FSDP wrapper in Fabric [converts the input path](https://github.com/Lightning-AI/pytorch-lightning/blob/0c9d4147e3569e24bf6130e78ae21b136562766e/src/lightning/fabric/strategies/fsdp.py#L440) to a pathlib.Path object. This conversion changes a valid URL like s3://bucket/xxx into an invalid format like s3:/bucket/xxx, effectively breaking the path.
To resolve this issue, path manipulation should be handled using os.path instead—[as done in TorchCheckpointIO](https://github.com/Lightning-AI/pytorch-lightning/blob/0c9d4147e3569e24bf6130e78ae21b136562766e/src/lightning/fabric/plugins/io/torch_io.py#L57).
### What version are you seeing the problem on?
v2.5
### How to reproduce the bug
```python
- Use FSDP as the strategy for fabric.
- Save a checkpoint with `fabric.save("s3://xxx")`
```
### Error messages and logs
```
# Error messages and logs here please
```
### Environment
_No response_
### More info
_No response_
cc @lantiga @justusschock
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/fabric/strategies/fsdp.py at the checkpoint path conversion around line 440, then compare the path handling in src/lightning/fabric/plugins/io/torch_io.py around line 57. Reproduce the issue with FSDP and fabric.save("s3://xxx"); done means S3 URLs remain valid for checkpoint save and load.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- cloud, distributed-systems
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 67/100