Lightning-AI / Lightning-AI/pytorch-lightning
Support memory snapshotting on OOM error
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Improve the user experience around out of memory errors
### Pitch
Support the following
```python
def oom_observer(device, alloc, device_alloc, device_free):
# snapshot right after an OOM happened
print('saving allocated state during OOM')
snapshot = torch.cuda.memory._snapshot()
dump(snapshot, open('oom_snapshot.pickle', 'wb'))
torch._C._cuda_attach_out_of_memory_observer(oom_observer)
```
Probably through the CUDAAccelerator, or an associated utility.
This would also require extra utilities to interpret the snapshot
### Alternatives
Not do it
### Additional context
Seen in https://zdevito.github.io/2022/08/16/memory-snapshots.html, https://zdevito.github.io/2022/12/09/memory-traces.html
cc @borda @carmocca @justusschock @awaelchli
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
The issue names the CUDAAccelerator, the out-of-memory observer API, and snapshot utilities but no files or tests. Start by locating existing CUDA out-of-memory handling and reviewing the proposed observer and snapshot usage. Clarify the required snapshot-interpretation utilities; done means snapshots are captured on OOM and can be interpreted as specified.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100