Add OOM observer with memory visualizations

[cli99]

This PR add a callback to generate visualizations of the state of allocated memory during an OutOfMemory exception. This callback registers an observer with the allocator that will be called everytime it is about to raise an OutOfMemoryError before any memory has been release while unwinding the exception. OOMObserver is attached to the Trainer at init stage.

To enable this, in the yaml config, do below, also need this change in the foundry mosaicml/llm-foundry#932

callbacks:
    oom_observer:
        {
            "folder": "traces",
            "overwrite": true,
            "filename": "rank{rank}_oom",
            "remote_file_name": "oci://bucket_name/{run_name}/oom_traces/rank{rank}_oom",
        }

The visualizations include

• snapshot of the memory state, filename_snapshot.pickle, can be used later for analysis with https://github.com/pytorch/pytorch/blob/main/torch/cuda/_memory_viz.py

• trace plot, filename_trace_plot.html
  

• segment plot, filename_segment_plot.html
  

• segment flamegraph, filename_segment_flamegraph.svg
  

• memory flamegrap, filename_memory_flamegraph.svg`
  

[j316chuck]

Thank you! This is much needed!

[j316chuck]

@cli99 mind adding a unit test to make sure we can add both a MemorySnapshot and an OOMSnapshot callback. Just want to test the scenario where someone calls torch.cuda._memory.snapshot twice.

[j316chuck]

Looks great @cli99 barring a couple of nits and the test case above 🎉 !

Also if you could figure out how to mock/add the unit test Mihir described that would be awesome but imo not blocking

[mvpatel2000]

LGTM, just some minor concerns :)

[cli99]

@cli99 mind adding a unit test to make sure we can add both a MemorySnapshot and an OOMSnapshot callback. Just want to test the scenario where someone calls torch.cuda._memory.snapshot twice.

added a test to show the two callbacks are compatible. Both torch.cuda._memory.snapshot and torch.cuda.memory._record_memory_history allows multiple calls. @j316chuck

[cli99]

@j316chuck, using warning.warn causes issues in a bunch of cpu callback tests, where they don't catch user warnings, so I switch to log.warning

[mvpatel2000]

The test is really nice!

I have some minor clean up comments but PR LGTM

[mvpatel2000]

LGTM!

Sorry for being a bit nit heavy on reviews 😂