add memory snapshot callback

[cli99]

This PR adds a callback to capture memory snapshot. A html file with interactive visualization would be generated if enabled. https://github.com/pytorch/pytorch.github.io/blob/site/assets/images/understanding-gpu-memory-1/snapshot.html shows an example visualization.

To enable the callback, in the yaml config file, also need this change in the foundry mosaicml/llm-foundry#810.

callbacks:
    memory_snapshot:
        {
            "skip_batches": 2,
            "interval": 3,
            "folder": "traces",
            "overwrite": true,
            "filename": "rank{rank}.{batch}.pt.trace.memory_snapshot.html",
            "remote_file_name": "oci://bucket_name/{run_name}/torch_memory_snapshots/rank{rank}.{batch}.pt.trace.memory_snapshot.html",
        }

Below is an example memory snapshot over three batches of MPT-7B with micro batch size 4, FSDP full shard, on 8XH100

For more details on the memory snapshot, refer to

• https://pytorch.org/blog/understanding-gpu-memory-1/#capturing-memory-snapshots
• https://pytorch.org/docs/main/torch_cuda_memory.html#generating-a-snapshot

[j316chuck]

Thanks for adding this Cheng! Overall, this PR looks good to me barring some style suggestions.

Additionally, here are two higher level use cases to consider for this PR:

• First, should we add this callback in the profiler instead so we support having N memory traces instead of just 1 memory trace? I can see pros and cons for both so just wanted to make sure we are aligned on this decision cc: @mvpatel2000
• Second, is it possible to add this memory trace when a run oom's?

Here's the original ticket that this PR addresses that has some ideas on the second point:
https://databricks.atlassian.net/browse/GRT-2231

[cli99]

Thanks for adding this Cheng! Overall, this PR looks good to me barring some style suggestions.

Additionally, here are two higher level use cases to consider for this PR:

• First, should we add this callback in the profiler instead so we support having N memory traces instead of just 1 memory trace? I can see pros and cons for both so just wanted to make sure we are aligned on this decision cc: @mvpatel2000
• Second, is it possible to add this memory trace when a run oom's?

Here's the original ticket that this PR addresses that has some ideas on the second point: https://databricks.atlassian.net/browse/GRT-2231

My thoughts on the two questions

1. Composer profiler aligns with what torch profiler supports and takes the schedule parameters; memory snapshot call back does not reply on any torch profiler apis (nor its schedule) and shall live by itself.
2. Generating memory snapshot right before OOM is a separate feature. It's going to generate a snapshot at a single timestamp (not over a time interval) and the best visualization i think would be a flame graph. Two options here: a) add to the current memory monitor callback(in its init function) b) implement as a separate callback. What do you think? @j316chuck @mvpatel2000

[mvpatel2000]

LGTM. Can you please provide a manual test and a screenshot just to show it works on a real run?

1. I agree this should be separate from profiler. It's convenient to have modularity, especially because profiler slows things down a lot
2. I think merging OOM capture into this callback is reasonable given it's the same torch system. I suggest follow-on PR since this is basically done

[j316chuck]

Agreed with 1 and 2 from Mihir. Up to you whether or not you want to create follow up PR or add the OOM callback to this PR 👍