Do memory analysis from Trace files

[payoto]

What does this PR do?

This PR makes it possible to analyse memory events from the Trace files to allow fine-grained analysis of memory allocations and to relate them to other events.

At the moment the memory timelines can be processed from the Pytorch profiler's export_memory_timeline. This output classifies the memory events in one of 7 categories during the export, however the categories are poorly documented and can be incorrect for more complex applications that use Pytorch distributed and torch.compile.

Instead this loads events in the trace dataframes:

  {
    "ph": "i", "cat": "cpu_instant_event", "s": "t", "name": "[memory]",
    "pid": 1401543, "tid": 1405042,
    "ts": 1336304245461.149,
    "args": {
      "Total Reserved": 82514542592, "Total Allocated": 61060382720, "Bytes": 8405385216, "Addr": 128803317743616, "Device Id": 0, "Device Type": 1, "Ev Idx": 6629
    }
  },

Sample usage of the feature - detailed usage in memory_analysis_demo_2.ipynb:

from hta.trace_analysis import TraceAnalysis
analyzer = TraceAnalysis(trace_dir=f"{trace_prefix}/tests/data/torchtitan_h100/")
memory_events = analyzer.get_memory_timeline()
# User specified categories:

def classify_big_allocs(row):
    if row["bytes_delta"] > 5e9:
        return "very_big_alloc (>5gB)"
    elif row["bytes_delta"] > 1e9:
        return "big_alloc (>1gB)"
    elif row["bytes_delta"] <1e6:
        return "small_alloc (<1mb)"
    return "normal_alloc"

alloc_sizes, _ = analyzer.get_memory_timeline_per_category(classification_func=classify_big_allocs)

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • Generate memory timelines from trace files #222
• Did you write any new necessary tests?
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you update the changelog?
  • N/A

[briancoutinho]

qq we are just turning the "dur" to 0 here , how does that keep the memoy events, i thought they are already retained

[payoto]

Memory events have no duration as they are "instantaneous" - as a consequence "dur" is set to na but the following lines do:

    # drop rows with null values
    df.dropna(axis=0, subset=["dur", "cat"], inplace=True)

Which dropped them.

[briancoutinho]

I think ev_idx is already there above see "index::ev_idx"