diff --git a/superbench/benchmarks/model_benchmarks/pytorch_base.py b/superbench/benchmarks/model_benchmarks/pytorch_base.py index de06b35d0..7810b982d 100644 --- a/superbench/benchmarks/model_benchmarks/pytorch_base.py +++ b/superbench/benchmarks/model_benchmarks/pytorch_base.py @@ -591,8 +591,11 @@ def _benchmark(self): Run the benchmark then handle post-run model log save/compare. Set SB_ENABLE_PYTORCH_PROFILER='1' to enable profiling. """ - # Check if this is a Nvidia GPU - if not (torch.cuda.is_available() and torch.version.cuda is not None): + # Check if this is a Nvidia or AMD GPU + if not ( + torch.cuda.is_available() and + (getattr(torch.version, 'cuda', None) is not None or getattr(torch.version, 'hip', None) is not None) + ): ok = super()._benchmark() self._post_run_model_log() return ok @@ -620,6 +623,9 @@ def _benchmark(self): local_rank = self._local_rank diag_agent_prof = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) + # Note: on ROCm builds of PyTorch, ProfilerActivity.CUDA and event.cuda_time are aliases that cover + # HIP kernels; the same code path therefore works for both NVIDIA and AMD GPUs (verified on MI300x, + # ROCm 6.3.4). dump_file_dir = os.environ.get('SB_TORCH_PROFILER_TRACE_DIR', '.') diag_agent_dump_file_path = f'{dump_file_dir}/torch-profiler-sb-{self._name}-{local_rank}.json' diag_agent_prof.__enter__() diff --git a/superbench/runner/runner.py b/superbench/runner/runner.py index a5ac13cbb..7a91abae6 100644 --- a/superbench/runner/runner.py +++ b/superbench/runner/runner.py @@ -7,6 +7,7 @@ import sys import json import random +import shlex import signal from pathlib import Path from pprint import pformat @@ -25,6 +26,17 @@ AnsibleClient = LazyImport('superbench.runner.ansible', 'AnsibleClient') +def _quote_for_bash_lc(value): + r"""Quote a value so it is safe both for the shell and for embedding inside an outer bash -lc '...' string. + + ``shlex.quote`` wraps values containing whitespace or shell metacharacters in single quotes. When the + resulting command is later interpolated into ``bash -lc '{command}'``, those single quotes would + terminate the outer single-quoted context. Escape any single quotes as ``'\''`` so the value survives + both quoting layers. + """ + return shlex.quote(value).replace("'", "'\\''") + + class SuperBenchRunner(): """SuperBench runner class.""" def __init__(self, sb_config, docker_config, ansible_config, sb_output_dir): @@ -112,6 +124,31 @@ def __get_enabled_benchmarks(self): return list(self._sb_config.superbench.enable) return [k for k, v in self._sb_benchmarks.items() if 'enable' in v and v.enable] + def __build_trace_command(self, benchmark_name, suffix, enable_nsys, enable_rocprof, trace_dir, rocprof_trace_dir): + """Build the profiler prefix (nsys or rocprofv2) to prepend to an execution command. + + Args: + benchmark_name (str): Benchmark name, used in the output filename. + suffix (str): Suffix to append to the output filename (e.g. rank id or empty string). + enable_nsys (bool): Whether nsys profiling is enabled. + enable_rocprof (bool): Whether rocprofv2 profiling is enabled. + trace_dir (str): Output directory for nsys traces. + rocprof_trace_dir (str): Output directory for rocprofv2 traces. + + Return: + str: Profiler command prefix with a trailing space, or an empty string if no profiler is enabled. + """ + if enable_nsys: + trace_output = _quote_for_bash_lc(f'{trace_dir}/{benchmark_name}{suffix}_traces') + return ( + f'nsys profile --output {trace_output} ' + f'--backtrace none --sample none --force-overwrite true --cpuctxsw none --trace cuda,nvtx ' + ) + if enable_rocprof: + trace_output = _quote_for_bash_lc(f'{rocprof_trace_dir}/{benchmark_name}{suffix}_traces') + return (f'rocprofv2 --hip-trace --kernel-trace --plugin json ' f'-d {trace_output} -- ') + return '' + def __get_mode_command(self, benchmark_name, mode, timeout=None): """Get runner command for given mode. @@ -135,12 +172,25 @@ def __get_mode_command(self, benchmark_name, mode, timeout=None): enable_nsys = os.environ.get('SB_ENABLE_NSYS', '') == '1' trace_dir = os.environ.get('SB_NSYS_TRACE_DIR', self._sb_output_dir) + # Enable rocprofv2 profiling based on environment variable + enable_rocprof = os.environ.get('SB_ENABLE_ROCPROF', '') == '1' + rocprof_trace_dir = os.environ.get('SB_ROCPROF_TRACE_DIR', self._sb_output_dir) + + # SB_ENABLE_NSYS and SB_ENABLE_ROCPROF are mutually exclusive; nsys takes precedence when both are set. + if enable_nsys and enable_rocprof: + logger.warning( + 'Both SB_ENABLE_NSYS and SB_ENABLE_ROCPROF are set; using nsys and ignoring rocprofv2. ' + 'Unset SB_ENABLE_NSYS to enable rocprofv2 tracing on AMD hardware.' + ) + enable_rocprof = False + mode_command = exec_command if mode.name == 'local': - trace_command = ( - f'nsys profile --output {trace_dir}/{benchmark_name}_{mode.proc_rank}_traces ' - f'--backtrace none --sample none --force-overwrite true --cpuctxsw none --trace cuda,nvtx ' - ) if enable_nsys and mode.proc_rank == 0 else '' + trace_command = '' + if mode.proc_rank == 0: + trace_command = self.__build_trace_command( + benchmark_name, f'_{mode.proc_rank}', enable_nsys, enable_rocprof, trace_dir, rocprof_trace_dir + ) # Build the command parts, only including trace if it's not empty command_parts = [] prefix = mode.prefix.format(proc_rank=mode.proc_rank, proc_num=mode.proc_num) @@ -159,23 +209,21 @@ def __get_mode_command(self, benchmark_name, mode, timeout=None): '--nnodes=$NNODES --node_rank=$NODE_RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT ' ) - nsys_prefix = ( - f'nsys profile --output {trace_dir}/{benchmark_name}_traces ' - f'--backtrace none --sample none --force-overwrite true --cpuctxsw none --trace cuda,nvtx ' - ) if enable_nsys else '' + trace_prefix = self.__build_trace_command( + benchmark_name, '', enable_nsys, enable_rocprof, trace_dir, rocprof_trace_dir + ) mode_command = ( - f'{nsys_prefix}' + f'{trace_prefix}' f'torchrun' f' --no_python --nproc_per_node={mode.proc_num} {torch_dist_params}{exec_command}' f' superbench.benchmarks.{benchmark_name}.parameters.distributed_impl=ddp' f' superbench.benchmarks.{benchmark_name}.parameters.distributed_backend=nccl' ) elif mode.name == 'mpi': - trace_command = ( - f'nsys profile --output {trace_dir}/{benchmark_name}_{mode.proc_rank}_traces ' - f'--backtrace none --sample none --force-overwrite true --cpuctxsw none --trace cuda,nvtx ' - ) if enable_nsys else '' + trace_command = self.__build_trace_command( + benchmark_name, f'_{mode.proc_rank}', enable_nsys, enable_rocprof, trace_dir, rocprof_trace_dir + ) mode_command = ( '{trace} ' 'mpirun ' # use default OpenMPI in image