Lightning-AI / Lightning-AI/pytorch-lightning
Provide teardown APIs in Fabric
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🚀 Feature
Right now, Fabric does not have any high-level APIs to free memory, unwrap models, etc.
### Motivation
Give the user device/strategy agnostic functions to call in their loops and training workflows when control is needed. Automate what can be automated.
### Pitch
- A way to unwrap layers that were wrapped with SyncBatchNorm (see also #13748)
- A way to free memory occupied by objects held within strategy
On the raw strategy/accelerator API level:
```py
accelerator.teadown_device()
model = strategy.teadown_module(model)
```
In the Fabric class
```py
fabric.free_memory()
model = fabric.teadown_module(model)
```
### Alternatives
- Have these methods only defined in the strategies.
- Keep the general purpose teardown() concept from PL which does many things, at the cost of flexibility in Lite.
### Additional context
Discussion started in https://github.com/Lightning-AI/lightning/pull/14662#discussion_r968623059
______________________________________________________________________
#### If you enjoy Lightning, check out our other projects! ⚡
- [**Metrics**](https://github.com/Lightning-AI/metrics): Machine learning metrics for distributed, scalable PyTorch applications.
- [**Lite**](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.html): enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.
- [**Flash**](https://github.com/Lightning-AI/lightning-flash): The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.
- [**Bolts**](https://github.com/Lightning-AI/lightning-bolts): Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.
- [**Lightning Transformers**](https://github.com/Lightning-AI/lightning-transformers): Flexible interface for high-performance research using SOTA Transformers leveraging PyTorch Lightning, Transformers, and Hydra.
cc @borda @tchaton @justusschock @awaelchli @carmocca
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 by reviewing the Fabric and raw strategy/accelerator APIs, then read the related SyncBatchNorm issue #13748 and the discussion linked from pull request #14662. Done would require agreed device-agnostic APIs for freeing strategy-held memory and unwrapping modules, with behavior validated across the relevant strategies.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100