Lightning-AI / Lightning-AI/pytorch-lightning

Provide teardown APIs in Fabric

Open
#14,682 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

design fabric feature
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.