Lightning-AI / Lightning-AI/pytorch-lightning

FSDP full state dict mangles fsspec path

Open
#20,406 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug ver: 2.4.x ver: 2.5.x waiting on author
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.