From 5d44815d95beedc90d097c40fbabb4a732ef2413 Mon Sep 17 00:00:00 2001 From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com> Date: Thu, 15 Feb 2024 07:48:45 -0700 Subject: [PATCH] Npkit updates (#1084) * removed warmup runs to be an opt in [ROCm/rccl commit: 16d7f372b7dd175181287159f3f79c3360746831] --- projects/rccl/tools/scripts/npkit_trace_generator.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/projects/rccl/tools/scripts/npkit_trace_generator.py b/projects/rccl/tools/scripts/npkit_trace_generator.py index 9ac472ad3c..d09b5b9ac0 100644 --- a/projects/rccl/tools/scripts/npkit_trace_generator.py +++ b/projects/rccl/tools/scripts/npkit_trace_generator.py @@ -55,7 +55,7 @@ def parse_cpu_event(event_bytes): 'timestamp': int.from_bytes(event_bytes[8:16], byteorder='little', signed=False) } -def parse_gpu_event_file_time(sync_dictionary, npkit_dump_dir, npkit_event_def, rank, buf_idx, gpu_clock_scale, cpu_clock_scale, dictionary_of_stats, warmup_runs=5): +def parse_gpu_event_file_time(sync_dictionary, npkit_dump_dir, npkit_event_def, rank, buf_idx, gpu_clock_scale, cpu_clock_scale, dictionary_of_stats, warmup_runs): gpu_event_file_path = os.path.join(npkit_dump_dir, 'gpu_events_rank_%d_buf_%d' % (rank, buf_idx)) stats_key = 'gpu_rank_%d' % (rank) channel_stats = {} @@ -286,7 +286,7 @@ def parse_cpu_event_file(npkit_dump_dir, npkit_event_def, rank, channel, cpu_clo -def convert_npkit_dump_to_trace(npkit_dump_dir, output_dir, npkit_event_def, gpu_statistics): +def convert_npkit_dump_to_trace(npkit_dump_dir, output_dir, npkit_event_def, gpu_statistics, warmup_runs=0): files_in_dump_dir = next(os.walk(npkit_dump_dir))[2] gpu_event_files = [x for x in files_in_dump_dir if x.startswith('gpu_events_rank_')] cpu_event_files = [x for x in files_in_dump_dir if x.startswith('cpu_events_rank_')] @@ -308,7 +308,7 @@ def convert_npkit_dump_to_trace(npkit_dump_dir, output_dir, npkit_event_def, gpu avg_time = {} number_events=0 for buf_idx in buf_indices: # get the avg time - parse_gpu_event_file_time(sync_dictionary, npkit_dump_dir, npkit_event_def, rank, buf_idx, gpu_clock_scale, cpu_clock_scale, dictionary_of_stats) + parse_gpu_event_file_time(sync_dictionary, npkit_dump_dir, npkit_event_def, rank, buf_idx, gpu_clock_scale, cpu_clock_scale, dictionary_of_stats, warmup_runs) for key in sync_dictionary: avg_time[key] = 0 @@ -317,7 +317,7 @@ def convert_npkit_dump_to_trace(npkit_dump_dir, output_dir, npkit_event_def, gpu avg_time[key] = avg_time[key] + (event['timestamp']/number_events) for buf_idx in buf_indices: - gpu_events = parse_gpu_event_file(avg_time, npkit_dump_dir, npkit_event_def, rank, buf_idx, gpu_clock_scale, cpu_clock_scale, dictionary_of_stats) + gpu_events = parse_gpu_event_file(avg_time, npkit_dump_dir, npkit_event_def, rank, buf_idx, gpu_clock_scale, cpu_clock_scale, dictionary_of_stats, warmup_runs) trace['traceEvents'].extend(gpu_events) @@ -345,9 +345,10 @@ if __name__ == '__main__': parser.add_argument('--npkit_event_header_path', type=str, required=True, help='Path to npkit_event.h.') parser.add_argument('--output_dir', type=str, required=True, help='Path to output directory.') parser.add_argument('--gpu_run_stats', type=bool, nargs='?', const=True, default=False, help="print stats instead.") + parser.add_argument('--warmup_runs', type=int, required=False, default=0, help="amount of warmup_runs on rccl.") args = parser.parse_args() gpu_statistics = False if args.gpu_run_stats is not None: gpu_statistics = args.gpu_run_stats npkit_event_def = parse_npkit_event_header(args.npkit_event_header_path) - convert_npkit_dump_to_trace(args.npkit_dump_dir, args.output_dir, npkit_event_def, gpu_statistics) \ No newline at end of file + convert_npkit_dump_to_trace(args.npkit_dump_dir, args.output_dir, npkit_event_def, gpu_statistics, args.warmup_runs) \ No newline at end of file