Lightning-AI / Lightning-AI/pytorch-lightning
Support model size calculation for FSDP & DeepSpeed (stage-3)
Open
Nobody has claimed this yet.
feature
strategy: deepspeed
strategy: fsdp
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🚀 Feature
Currently `LightningModule.model_size` (deprecated) relies on:
https://github.com/PyTorchLightning/pytorch-lightning/blob/45c45dc7b018f9a2db60f5df1a3f7dbbb45ccb36/pytorch_lightning/utilities/memory.py#L165-L179
to calculate model size. Since we are going to keep `get_model_size_mb` utility function, we should have a way to calculate model_sizes for sharded models.
cc @borda @awaelchli @carmocca @SeanNaren @rohitgr7 @akihironitta
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.
Assessment
This issue has not been assessed yet.