Lightning-AI / Lightning-AI/pytorch-lightning
Use new state-dict APIs in FSDPStrategy
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
In PyTorch 2.4, the [state-dict context managers we use in FSDPStrategy](https://github.com/Lightning-AI/pytorch-lightning/blob/5829ef8ab3bfa3eb03cfc35e842bef6ebd6bf007/src/lightning/fabric/strategies/fsdp.py#L800-L831) are being deprecated with FutureWarning in PyTorch 2.4. In #20010 I'm suppressing these warnings temporarily.
### Pitch
Use the new `torch.distributed.checkpoint` APIs for PyTorch >= 2.4.
### Alternatives
_No response_
### Additional context
_No response_
cc @borda @justusschock @awaelchli @carmocca
Contributor guide
Research direction
Start in src/lightning/fabric/strategies/fsdp.py around the state-dict context managers at lines 800-831. Review the torch.distributed.checkpoint APIs for PyTorch >=2.4 and replace the deprecated usage; the work is done when those APIs are used without the FutureWarning described in the issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100