Zizeng Meng 861945100e [Kineto] Enable OOM observer (#152160) ...

Summary:
# Context:
When memory leak happens, it usually trigger the OOM in the later iterations. The snapshot of full iteration will be huge and hard to interpret.
On CUDA side, they provide OOM observer which generates snapshot when OOM happens with latest 1,500,000 entries for debugging.

In this diff, we want to implement the feature on MTIA side

Test Plan:
Run this test with last diff in the stack.
```
buck run @//mode/opt  kineto/libkineto/fb/mtia/integration_tests:mtia_memory_auto_trace_test
```

As shown, the memory_snapshot is generated when oom happens
Log: P1794792326
Snapshot: https://fburl.com/pytorch_memory_visualizer/lx73y6s3 {F1977402355}

Differential Revision: D71993315

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152160
Approved by: https://github.com/sraikund16

2025-04-27 15:56:44 +00:00

..

__init__.py

[Kineto] Enable OOM observer (#152160)

2025-04-27 15:56:44 +00:00

_utils.py

torch.mtia module for MTIA device backend (#123612)

2024-04-26 16:17:54 +00:00

memory.py

[MTIA] (4/n) Implement PyTorch APIs to query/reset device peak memory usage (#146751)

2025-02-11 03:51:48 +00:00