Lightning-AI / Lightning-AI/lightning-thunder
[FSDP] Support optimizer state checkpointing
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
### Motivation
Saving the optimizer state is critical to resume a training run.
### Pitch
```python
from thunder.distributed.checkpoint import get_optimizer_state_dict, load_optimizer_state_dict
```
from https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/distributed/checkpoint.py
Then, integrate it into Fabric's Thunder FSDP strategy
### References
`get_optimizer_state_dict`: https://github.com/pytorch/pytorch/blob/ee557d8f61bbe8a54742a82507a6edb2de3e5a89/torch/distributed/checkpoint/state_dict.py#L662
`_optim_state_dict`: https://github.com/pytorch/pytorch/blob/ee557d8f61bbe8a54742a82507a6edb2de3e5a89/torch/distributed/fsdp/_optim_utils.py#L1864
### Additional context
Follow-up to Lightning-AI/lit-thunder-LEGACY#1909
cc @carmocca @awaelchli @crcrpar
Contributor guide
No contributing guide indexed for this repository
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 with thunder/distributed/checkpoint.py and read PyTorch's referenced get_optimizer_state_dict and _optim_state_dict implementations. Then trace Fabric's Thunder FSDP strategy entry points and determine how optimizer state saving and loading should be integrated. Done means optimizer state can be checkpointed and restored when resuming a Thunder FSDP training run.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100