From 2b9a9f314d77202b265c2d9ce18d74699ab3818d Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Tue, 9 Jul 2024 23:18:32 +0530 Subject: [PATCH] Add plot-metric arc --- neuralmagic/tools/profiler/visualize_trace.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/neuralmagic/tools/profiler/visualize_trace.py b/neuralmagic/tools/profiler/visualize_trace.py index b8ccdf612aca7..844a53c4ac5a3 100644 --- a/neuralmagic/tools/profiler/visualize_trace.py +++ b/neuralmagic/tools/profiler/visualize_trace.py @@ -223,16 +223,14 @@ def is_vocab_embedding_op(op_name: str): def plot_trace_df(traces_df: pd.DataFrame, + plot_metric: str, plot_title: str, output: Optional[Path] = None): - traces_df["cuda_time_ms"] = traces_df["cuda_time_us"] / 1000 - traces_df = traces_df.fillna(0) - phases = traces_df['phase'].unique() traces_df = traces_df.pivot_table(index="phase", columns="name", - values="cuda_time_ms", + values=plot_metric, aggfunc="sum") traces_df = group_trace_by_operations(traces_df) @@ -281,6 +279,7 @@ def main( json_trace: Path, output_directory: Path, depth: int, # Fetch/Plot operations at this depth of the Json tree + plot_metric: str, make_names_unique: bool, top_k: int, json_nodes_to_fold: List[str]): @@ -329,6 +328,11 @@ def keep_only_top_entries(df: pd.DataFrame, # Combine all data frames so they can be put in a single plot traces_df = pd.concat(trace_dfs) + + # Add a derived metric `cuda_time_ms` + traces_df["cuda_time_ms"] = traces_df["cuda_time_us"] / 1000 + traces_df = traces_df.fillna(0) + return traces_df def make_plot_title_suffix(profile_json: dict) -> str: @@ -365,9 +369,9 @@ def make_plot_title_suffix(profile_json: dict) -> str: plot_title_suffix = make_plot_title_suffix(profile_json) - plot_trace_df(prefill_traces, "prefill " + plot_title_suffix, + plot_trace_df(prefill_traces, plot_metric, "prefill " + plot_title_suffix, output_directory / Path("prefill.png")) - plot_trace_df(decode_traces, "decodes " + plot_title_suffix, + plot_trace_df(decode_traces, plot_metric, "decodes " + plot_title_suffix, output_directory / Path("decode_steps.png")) @@ -394,6 +398,10 @@ def make_plot_title_suffix(profile_json: dict) -> str: parser.add_argument("--fold-json-node", nargs='+', default=['Sampler', 'LogitsProcessor']) + parser.add_argument("--plot-metric", + type=str, + default="cuda_time_ms", + help='Metric to plot. some options are cuda_time_us, cuda_time_ms, pct_cuda_time') parser.add_argument( "--step-plot-interval", type=int, @@ -415,5 +423,5 @@ def make_plot_title_suffix(profile_json: dict) -> str: output_directory = args.output_directory if args.output_directory else Path( args.json_trace).parent - main(Path(args.json_trace), output_directory, depth, make_names_unique, - args.top_k, args.fold_json_node) + main(Path(args.json_trace), output_directory, depth, + args.plot_metric, make_names_unique, args.top_k, args.fold_json_node)