Open
Description
New Feature
We can use memory instant events in pytorch profiler result to generate a nice gpu memory trace in perfetto:
This memory trace is aligned to python stack trace so it's more useful than pytorch solutions such as export_memory_timeline
I have implemented this by add profile_memory to torch.profiler.profile and parse pytorch trace events, I think viztracer can add this feature easily. we may use --log_torch and --log_torch_memory to enable this feature.
for ev in pth_trace_dict["traceEvents"]:
if (ev["ph"] == "I" or ev["ph"] == "i") and "args" in ev:
if "Total Allocated" in ev["args"]:
total_alloc = ev["args"]["Total Allocated"]
total_reserv = ev["args"]["Total Reserved"]
if ev["args"]["Device Id"] >= 0:
external_events.append({
"name": "pth_memory",
"ph": "C",
"pid": pid,
"tid": 0,
"ts": ev["ts"] + base_ts_us,
"args": {
"Total Allocated": total_alloc,
"Total Reserved": total_reserv,
},
})