Lightning-AI / Lightning-AI/pytorch-lightning
FSDP full state dict mangles fsspec path
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
In `FSDPStrategy.save_checkpoint`, the `filepath` variable is transformed via
https://github.com/Lightning-AI/pytorch-lightning/blob/3627c5bfac704d44c0d055a2cdf6f3f9e3f9e8c1/src/lightning/pytorch/strategies/fsdp.py#L562
This only makes sense if doing sharded checkpointing, and in fact mangles any legitimate fsspec path that is passed in.
When `self._state_dict_type == "full"`,
```
super().save_checkpoint(checkpoint=checkpoint, filepath=path)
```
is called, using the normal CheckpointIO workflow, but with the mangled path.
The expected behavior should be that if the user chooses full state dict type, CheckpointIO and remote paths should work as usual, but currently full state dict checkpoints cannot be saved to remote paths.
### What version are you seeing the problem on?
v2.4
### How to reproduce the bug
```python
trainer = L.Trainer(
strategy="fsdp"
default_root_dir="s3://example/path"
)
trainer.fit(model=...)
```
### Error messages and logs
```
# Error messages and logs here please
```
### Environment
Current environment
```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```
### More info
_No response_
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 at src/lightning/pytorch/strategies/fsdp.py around line 562 and trace how filepath is transformed before the full-state branch calls CheckpointIO. Reproduce with the S3-style default_root_dir shown in the issue; done means full state dict checkpoints preserve legitimate fsspec paths and save through the normal CheckpointIO workflow.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- cloud, distributed-systems
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100