Skip to content

Commit 8e853d5

Browse files
lihongyang1990lihongyang1990
committed
Support hipprof step profiling
1 parent 75ad1ee commit 8e853d5

3 files changed

Lines changed: 199 additions & 7 deletions

File tree

‎flagscale/train/megatron/training/config/common_config.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,15 @@ class ProfilingConfig:
4141
use_pytorch_profiler: bool = False
4242
"""Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard."""
4343

44+
use_hipprof_profiler: bool = False
45+
"""Use hipprof session control for profiling."""
46+
47+
hipprof_bin_path: str = ""
48+
"""Path to the hipprof executable used to send session start/stop commands."""
49+
50+
hipprof_session_id: str = ""
51+
"""hipprof session id to control. If empty, HIPPROF_SESSION_ID is used."""
52+
4453
pytorch_profiler_collect_shapes: bool = False
4554
"""Collect tensor shape in pytorch profiler."""
4655

‎flagscale/train/megatron/training/training.py‎

Lines changed: 70 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ def set_startup_timestamps(program_start=None, main_entry=None):
4343
import math
4444
import os
4545
import socket
46+
import shlex
47+
import subprocess
4648
import sys
4749
from contextlib import nullcontext
4850
from pathlib import Path
@@ -344,6 +346,60 @@ def _is_global_rank_zero() -> bool:
344346
return True
345347

346348

349+
def _get_global_rank() -> int:
350+
"""Return global distributed rank, defaulting to 0 before distributed init."""
351+
if torch.distributed.is_available() and torch.distributed.is_initialized():
352+
return torch.distributed.get_rank()
353+
return 0
354+
355+
356+
def _is_profile_rank(args) -> bool:
357+
return len(args.profile_ranks) == 0 or _get_global_rank() in args.profile_ranks
358+
359+
360+
def _format_command(cmd):
361+
return " ".join(shlex.quote(str(part)) for part in cmd)
362+
363+
364+
def _hipprof_session_control(args, action):
365+
if action not in ("start", "stop"):
366+
raise ValueError(f"Unsupported hipprof session action: {action}")
367+
368+
session_id = getattr(args, "hipprof_session_id", "") or os.environ.get("HIPPROF_SESSION_ID", "")
369+
if not session_id:
370+
raise RuntimeError(
371+
"hipprof profiling requires train.model.hipprof_session_id "
372+
"or HIPPROF_SESSION_ID."
373+
)
374+
375+
hipprof_bin = (
376+
getattr(args, "hipprof_bin_path", "")
377+
or os.environ.get("HIPPROF_BIN_PATH", "")
378+
or "hipprof"
379+
)
380+
cmd = [hipprof_bin, "--session", session_id, f"--{action}"]
381+
try:
382+
subprocess.run(cmd, check=True, capture_output=True, text=True)
383+
except FileNotFoundError as exc:
384+
raise RuntimeError(
385+
f"hipprof executable not found: {hipprof_bin}. "
386+
"Set train.model.hipprof_bin_path or HIPPROF_BIN_PATH."
387+
) from exc
388+
except subprocess.CalledProcessError as exc:
389+
output = []
390+
if exc.stdout:
391+
output.append(f"stdout:\n{exc.stdout.strip()}")
392+
if exc.stderr:
393+
output.append(f"stderr:\n{exc.stderr.strip()}")
394+
details = "\n".join(output)
395+
if details:
396+
details = "\n" + details
397+
raise RuntimeError(
398+
f"hipprof session {action} failed with exit code {exc.returncode}: "
399+
f"{_format_command(cmd)}{details}"
400+
) from exc
401+
402+
347403
from megatron.core.msc_utils import MultiStorageClientFeature, open_file
348404

349405

@@ -2821,14 +2877,15 @@ def post_training_step_callbacks(
28212877
if (
28222878
args.profile
28232879
and iteration == args.profile_step_end
2824-
and (len(args.profile_ranks) == 0 or
2825-
torch.distributed.get_rank() in args.profile_ranks)
2880+
and _is_profile_rank(args)
28262881
):
28272882
if args.use_pytorch_profiler:
28282883
assert prof is not None
28292884
prof.stop()
28302885
if prof.execution_trace_observer is not None:
28312886
prof.execution_trace_observer.unregister_callback()
2887+
elif getattr(args, "use_hipprof_profiler", False):
2888+
_hipprof_session_control(args, "stop")
28322889
else:
28332890
torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStop())
28342891
if nsys_nvtx_context is not None:
@@ -3204,8 +3261,13 @@ def get_e2e_base_metrics():
32043261
nsys_nvtx_context = None # reference to context for nsys profiling, so it can be cleaned up
32053262
if (
32063263
args.profile
3207-
and (len(args.profile_ranks) == 0 or
3208-
torch.distributed.get_rank() in args.profile_ranks)
3264+
and args.use_pytorch_profiler
3265+
and getattr(args, "use_hipprof_profiler", False)
3266+
):
3267+
raise RuntimeError("use_pytorch_profiler and use_hipprof_profiler cannot be enabled together.")
3268+
if (
3269+
args.profile
3270+
and _is_profile_rank(args)
32093271
and args.use_pytorch_profiler
32103272
):
32113273
if args.pytorch_profiler_collect_chakra:
@@ -3264,11 +3326,12 @@ def trace_handler(p):
32643326
# Run training iterations till done.
32653327
buffered_rollouts = None
32663328
while iteration < args.train_iters:
3267-
if (args.profile
3268-
and (len(args.profile_ranks) == 0 or
3269-
torch.distributed.get_rank() in args.profile_ranks)):
3329+
if args.profile and _is_profile_rank(args):
32703330
if args.use_pytorch_profiler:
32713331
prof.step()
3332+
elif getattr(args, "use_hipprof_profiler", False):
3333+
if iteration == args.profile_step_start:
3334+
_hipprof_session_control(args, "start")
32723335
elif iteration == args.profile_step_start:
32733336
torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStart())
32743337
nsys_nvtx_context = torch.autograd.profiler.emit_nvtx(record_shapes=True)

‎tools/profiling/hipprof_python_wrapper.sh‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
#!/usr/bin/env bash
2+
set -euo pipefail
3+
4+
real_python="${HIPPROF_REAL_PYTHON:-}"
5+
if [[ -z "$real_python" ]]; then
6+
real_python="$(command -v python3 || command -v python)"
7+
fi
8+
9+
infer_cli_values() {
10+
local flag="$1"
11+
shift
12+
local collecting=0
13+
local values=()
14+
for arg in "$@"; do
15+
if [[ "$arg" == "$flag" ]]; then
16+
collecting=1
17+
continue
18+
fi
19+
if [[ "$collecting" == "1" ]]; then
20+
if [[ "$arg" == --* ]]; then
21+
break
22+
fi
23+
values+=("$arg")
24+
fi
25+
done
26+
if [[ "${#values[@]}" -gt 0 ]]; then
27+
local IFS=,
28+
echo "${values[*]}"
29+
fi
30+
}
31+
32+
global_rank="${RANK:-na}"
33+
profile_ranks="${HIPPROF_PROFILE_RANKS:-}"
34+
if [[ -z "$profile_ranks" ]]; then
35+
profile_ranks="$(infer_cli_values --profile-ranks "$@")"
36+
fi
37+
profile_ranks="${profile_ranks// /}"
38+
profile_ranks="${profile_ranks#[}"
39+
profile_ranks="${profile_ranks%]}"
40+
if [[ -n "$profile_ranks" && "$global_rank" != "na" ]]; then
41+
case ",${profile_ranks}," in
42+
*,"${global_rank}",*) ;;
43+
*) exec "$real_python" "$@" ;;
44+
esac
45+
fi
46+
47+
hipprof_bin="${HIPPROF_BIN_PATH:-}"
48+
if [[ -z "$hipprof_bin" ]]; then
49+
hipprof_bin="$(infer_cli_values --hipprof-bin-path "$@")"
50+
fi
51+
hipprof_bin="${hipprof_bin:-hipprof}"
52+
output_root="${HIPPROF_OUTPUT_DIR:-}"
53+
if [[ -z "$output_root" ]]; then
54+
echo "HIPPROF_OUTPUT_DIR is required when using hipprof_python_wrapper.sh" >&2
55+
exit 2
56+
fi
57+
58+
host="${HOSTNAME:-$(hostname)}"
59+
local_rank="${LOCAL_RANK:-na}"
60+
session_prefix="${HIPPROF_SESSION_ID_PREFIX:-flagscale}"
61+
session_id="${HIPPROF_SESSION_ID:-${session_prefix}_${host}_rank${global_rank}_pid$$}"
62+
export HIPPROF_BIN_PATH="$hipprof_bin"
63+
export HIPPROF_SESSION_ID="$session_id"
64+
65+
out_dir="${output_root}/${host}_global${global_rank}_local${local_rank}_pid$$"
66+
mkdir -p "$out_dir"
67+
68+
trace_args=()
69+
api_trace="$(printf '%s%s%s' H S A)"
70+
trace_list="${HIPPROF_TRACE:-HIP,RCCL,${api_trace}}"
71+
IFS=',' read -r -a traces <<< "$trace_list"
72+
for trace in "${traces[@]}"; do
73+
trace_key="$(printf '%s' "${trace// /}" | tr '[:lower:]' '[:upper:]')"
74+
trace_arg="$(printf '%s' "$trace_key" | tr '[:upper:]' '[:lower:]')"
75+
if [[ "$trace_key" == "" || "$trace_key" == "NONE" ]]; then
76+
continue
77+
fi
78+
if [[ "$trace_key" == "HIP" || "$trace_key" == "RCCL" || "$trace_key" == "$api_trace" ]]; then
79+
trace_args+=("--${trace_arg}-trace")
80+
continue
81+
fi
82+
echo "Unsupported HIPPROF_TRACE entry: $trace" >&2
83+
exit 2
84+
done
85+
86+
group_stream="${HIPPROF_GROUP_STREAM:-1}"
87+
group_args=()
88+
if [[ "$group_stream" == "1" || "$group_stream" == "true" ]]; then
89+
group_args+=(--group-stream)
90+
fi
91+
92+
segment_size="${HIPPROF_SEGMENT_SIZE:-50000}"
93+
segment_args=()
94+
if [[ -n "$segment_size" ]]; then
95+
segment_args+=(--segment-size "$segment_size")
96+
fi
97+
98+
"$hipprof_bin" \
99+
--session "$session_id" \
100+
"${trace_args[@]}" \
101+
"${group_args[@]}" \
102+
"${segment_args[@]}" \
103+
-d "$out_dir" \
104+
-o "$out_dir/result" \
105+
"$real_python" "$@" &
106+
107+
child=$!
108+
cleanup() {
109+
"$hipprof_bin" --session "$session_id" --stop >/dev/null 2>&1 || true
110+
}
111+
trap cleanup EXIT
112+
113+
sleep "${HIPPROF_SESSION_WARMUP_SECONDS:-2}"
114+
cleanup
115+
116+
set +e
117+
wait "$child"
118+
status=$?
119+
set -e
120+
exit "$status"

0 commit comments

Comments
 (0)