From 3336aa191c3df3027a550f7ac87c1d98b31cf197 Mon Sep 17 00:00:00 2001 From: Anupam Bhatnagar Date: Mon, 21 Aug 2023 17:20:13 +0000 Subject: [PATCH] Adding allocated and reserved memory values to memory timline view. (#107056) Summary: This diff adds the max allocated and max reserved memory values to the memory timeline plot. Test Plan: Executed `buck run mode/dev-nosan kineto/libkineto/fb/integration_tests:pytorch_resnet_integration_test -- --enable_profiling --profile_memory --trace_handler=auto_trace --with_stack --record_shapes` on my devgpu. The generated output is at https://www.internalfb.com/manifold/explorer/ai_efficiency/tree/traces/dynocli/devgpu020.odn1.facebook.com/rank-0/rank-0.Aug_10_16_50_50.236946.pt.memorytl.html {F1067885545} Screenshot of the html above {F1067886350} Reviewed By: aaronenyeshi Differential Revision: D48251791 Pull Request resolved: https://github.com/pytorch/pytorch/pull/107056 Approved by: https://github.com/aaronenyeshi, https://github.com/davidberard98 --- torch/profiler/_memory_profiler.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/torch/profiler/_memory_profiler.py b/torch/profiler/_memory_profiler.py index 336f8d056e6cf6..19488ee7fe5077 100644 --- a/torch/profiler/_memory_profiler.py +++ b/torch/profiler/_memory_profiler.py @@ -1151,6 +1151,8 @@ def export_memory_timeline_html( mt = self._coalesce_timeline(device) times, sizes = np.array(mt[0]), np.array(mt[1]) stacked = np.cumsum(sizes, axis=1) / 1024**3 + max_memory_allocated = torch.cuda.max_memory_allocated() + max_memory_reserved = torch.cuda.max_memory_reserved() # Plot memory timeline as stacked data fig = plt.figure(figsize=figsize, dpi=80) @@ -1164,7 +1166,11 @@ def export_memory_timeline_html( axes.set_xlabel("Time (us)") axes.set_ylabel("Memory (GB)") title = "\n\n".join( - ([title] if title else []) + [f"Max: {stacked[:, -1].max():.2f} GB"] + ([title] if title else []) + + [ + f"Max memory allocated: {max_memory_allocated/(10**9):.2f} GB \n" + f"Max memory reserved: {max_memory_reserved/(10**9):.2f} GB" + ] ) axes.set_title(title)