Support hipprof step profiling

flagscale/runner/launcher/launcher_ssh.py

@@ -11,6 +11,10 @@
 
 from flagscale.runner.elastic.monitor_service import MonitorService
 from flagscale.runner.launcher.launcher_base import LauncherBase
+from flagscale.runner.profiling import (
+    configure_hipprof_env,
+    remove_launcher_profiling_args,
+)
 from flagscale.runner.utils import (
     JobStatus,
     add_decive_extra_config,
@@ -110,10 +114,7 @@ def _get_runner_cmd_train(
     if "enable_gpu_health_check" in runner_args:
         del runner_args["enable_gpu_health_check"]
     # Profiling options are consumed by launcher; torchrun doesn't accept them.
-    if "nsys_bin_path" in runner_args:
-        del runner_args["nsys_bin_path"]
-    if "nsys_rep_file_path" in runner_args:
-        del runner_args["nsys_rep_file_path"]
+    remove_launcher_profiling_args(runner_args)
     if "deploy" in runner_args:
         del runner_args["deploy"]
     if "enable_perf_monitor" in runner_args:
@@ -227,6 +228,13 @@ def _run_each(
         cur_envs=None,
         enable_monitoring=False,
     ):
+        if self.task_type == "train":
+            cur_envs = configure_hipprof_env(
+                self.config.experiment.runner,
+                self.config.train.model,
+                cur_envs,
+            )
+
         export_cmd = []
         if cur_envs:
             for k, v in cur_envs.items():

flagscale/runner/profiling.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+HIPPROF_RUNNER_KEYS = ("hipprof_bin_path", "hipprof_output_dir")
+HIPPROF_WRAPPER_PATH = "tools/profiling/hipprof_python_wrapper.sh"
+LAUNCHER_PROFILING_RUNNER_KEYS = (
+    "nsys_bin_path",
+    "nsys_rep_file_path",
+    *HIPPROF_RUNNER_KEYS,
+)
+
+
+def remove_launcher_profiling_args(runner_args: dict) -> None:
+    for key in LAUNCHER_PROFILING_RUNNER_KEYS:
+        runner_args.pop(key, None)
+
+
+def configure_hipprof_env(
+    runner_config: Mapping,
+    model_config: Mapping,
+    current_env: Mapping | None = None,
+) -> dict:
+    env = dict(current_env or {})
+    hipprof_bin_path = runner_config.get("hipprof_bin_path", None)
+    hipprof_output_dir = runner_config.get("hipprof_output_dir", None)
+    enabled = bool(model_config.get("use_hipprof_profiler", False))
+    requested = enabled or hipprof_bin_path is not None or hipprof_output_dir is not None
+
+    if not requested:
+        return env
+    if not enabled:
+        raise ValueError(
+            "hipprof runner configuration requires train.model.use_hipprof_profiler: true"
+        )
+    if not model_config.get("profile", False):
+        raise ValueError("hipprof step profiling requires train.model.profile: true")
+    if not hipprof_bin_path or not hipprof_output_dir:
+        raise ValueError(
+            "hipprof profiling requires both experiment.runner.hipprof_bin_path "
+            "and experiment.runner.hipprof_output_dir"
+        )
+    if runner_config.get("nsys_bin_path", None) or runner_config.get("nsys_rep_file_path", None):
+        raise ValueError("nsys and hipprof launcher profiling cannot be enabled together")
+
+    hipprof_executable = str(hipprof_bin_path)
+    if os.path.basename(hipprof_executable) != "hipprof":
+        hipprof_executable = os.path.join(hipprof_executable, "hipprof")
+
+    previous_python = env.get("PYTHON_EXEC")
+    if previous_python and not env.get("HIPPROF_REAL_PYTHON"):
+        env["HIPPROF_REAL_PYTHON"] = previous_python
+    env["PYTHON_EXEC"] = HIPPROF_WRAPPER_PATH
+    env["HIPPROF_BIN_PATH"] = hipprof_executable
+    env["HIPPROF_OUTPUT_DIR"] = str(hipprof_output_dir)
+    return env

flagscale/train/megatron/hipprof.py

@@ -0,0 +1,43 @@
+import os
+import shlex
+import subprocess
+
+
+def _format_command(command):
+    return " ".join(shlex.quote(str(part)) for part in command)
+
+
+def hipprof_session_control(action):
+    if action not in ("start", "stop"):
+        raise ValueError(f"Unsupported hipprof session action: {action}")
+
+    session_id = os.environ.get("HIPPROF_SESSION_ID", "")
+    if not session_id:
+        raise RuntimeError(
+            "hipprof profiling requires a launcher-provided session. "
+            "Configure experiment.runner.hipprof_bin_path and "
+            "experiment.runner.hipprof_output_dir."
+        )
+
+    hipprof_bin = os.environ.get("HIPPROF_BIN_PATH", "") or "hipprof"
+    command = [hipprof_bin, "--session-client", session_id, f"--{action}"]
+    try:
+        subprocess.run(command, check=True, capture_output=True, text=True)
+    except FileNotFoundError as error:
+        raise RuntimeError(
+            f"hipprof executable not found: {hipprof_bin}. Configure "
+            "experiment.runner.hipprof_bin_path."
+        ) from error
+    except subprocess.CalledProcessError as error:
+        output = []
+        if error.stdout:
+            output.append(f"stdout:\n{error.stdout.strip()}")
+        if error.stderr:
+            output.append(f"stderr:\n{error.stderr.strip()}")
+        details = "\n".join(output)
+        if details:
+            details = "\n" + details
+        raise RuntimeError(
+            f"hipprof session {action} failed with exit code {error.returncode}: "
+            f"{_format_command(command)}{details}"
+        ) from error

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

@@ -41,6 +41,11 @@ class ProfilingConfig:
     use_pytorch_profiler: bool = False
     """Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard."""
 
+    use_hipprof_profiler: bool = False
+    """Use hipprof session control for profiling. Configure the hipprof executable and output
+    directory under experiment.runner.
+    """
+
     pytorch_profiler_collect_shapes: bool = False
     """Collect tensor shape in pytorch profiler."""
 

flagscale/train/megatron/training/training.py

@@ -146,6 +146,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 )
 from megatron.core.fp8_utils import correct_amax_history_if_needed
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.hipprof import hipprof_session_control
 from megatron.core.pipeline_parallel.utils import (
     is_pp_first_stage,
     is_pp_last_stage,
@@ -344,6 +345,17 @@ def _is_global_rank_zero() -> bool:
     return True
 
 
+def _get_global_rank() -> int:
+    """Return global distributed rank, defaulting to 0 before distributed init."""
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    return 0
+
+
+def _is_profile_rank(args) -> bool:
+    return len(args.profile_ranks) == 0 or _get_global_rank() in args.profile_ranks
+
+
 from megatron.core.msc_utils import MultiStorageClientFeature, open_file
 
 
@@ -2821,14 +2833,15 @@ def post_training_step_callbacks(
     if (
         args.profile
         and iteration == args.profile_step_end
-        and (len(args.profile_ranks) == 0 or
-             torch.distributed.get_rank() in args.profile_ranks)
+        and _is_profile_rank(args)
     ):
         if args.use_pytorch_profiler:
             assert prof is not None
             prof.stop()
             if prof.execution_trace_observer is not None:
                 prof.execution_trace_observer.unregister_callback()
+        elif getattr(args, "use_hipprof_profiler", False):
+            hipprof_session_control("stop")
         else:
             torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStop())
             if nsys_nvtx_context is not None:
@@ -3204,8 +3217,13 @@ def get_e2e_base_metrics():
     nsys_nvtx_context = None # reference to context for nsys profiling, so it can be cleaned up
     if (
         args.profile
-        and (len(args.profile_ranks) == 0 or
-             torch.distributed.get_rank() in args.profile_ranks)
+        and args.use_pytorch_profiler
+        and getattr(args, "use_hipprof_profiler", False)
+    ):
+        raise RuntimeError("use_pytorch_profiler and use_hipprof_profiler cannot be enabled together.")
+    if (
+        args.profile
+        and _is_profile_rank(args)
         and args.use_pytorch_profiler
     ):
         if args.pytorch_profiler_collect_chakra:
@@ -3264,11 +3282,12 @@ def trace_handler(p):
     # Run training iterations till done.
     buffered_rollouts = None
     while iteration < args.train_iters:
-        if (args.profile 
-            and (len(args.profile_ranks) == 0 or
-                 torch.distributed.get_rank() in args.profile_ranks)):
+        if args.profile and _is_profile_rank(args):
             if args.use_pytorch_profiler:
                 prof.step()
+            elif getattr(args, "use_hipprof_profiler", False):
+                if iteration == args.profile_step_start:
+                    hipprof_session_control("start")
             elif iteration == args.profile_step_start:
                 torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStart())
                 nsys_nvtx_context = torch.autograd.profiler.emit_nvtx(record_shapes=True)

pyproject.toml

@@ -121,3 +121,7 @@ extend-ignore-identifiers-re = [
     "nd_*",
     "\\d+[Bb][-]?[Aa]\\d+(\\.\\d+)?[Bb]", # protect 30BA3B
 ]
+
+[tool.typos.default.extend-words]
+HSA = "HSA"
+hsa = "hsa"

tests/functional_tests/train/deepseek/conf/tp1_pp2_ep2_v4.yaml

@@ -14,6 +14,8 @@ experiment:
     nproc_per_node: 8
     # nsys_bin_path: /usr/local/bin/nsys
     # nsys_rep_file_path: ${experiment.exp_dir}/nsys_report/
+    # hipprof_bin_path: /opt/dtk/bin/hipprof
+    # hipprof_output_dir: ${experiment.exp_dir}/hipprof_report/
   shell_cmds: null
   envs:
     HYDRA_FULL_ERROR: 1

tests/functional_tests/train/deepseek/conf/train/tp1_pp2_ep2_v4.yaml

@@ -135,6 +135,13 @@ model:
   # tensorboard_dir: ${experiment.exp_dir}/tensorboard
   # pytorch_profiler_collect_callstack: true
 
+  # hipprof profile args ==============
+  # profile: true
+  # use_hipprof_profiler: true
+  # profile_step_start: 4
+  # profile_step_end: 5
+  # profile_ranks: [0]
+
   optimizer:
     optimizer: dist_muon
     muon_tp_mode: distributed

tests/unit_tests/runner/test_hipprof_profiling.py

@@ -0,0 +1,92 @@
+import unittest
+
+from flagscale.runner.profiling import (
+    configure_hipprof_env,
+    remove_launcher_profiling_args,
+)
+
+
+class HipprofEnvironmentTests(unittest.TestCase):
+    def setUp(self):
+        self.runner = {
+            "hipprof_bin_path": "/opt/dtk/bin",
+            "hipprof_output_dir": "/tmp/hipprof",
+        }
+        self.model = {
+            "profile": True,
+            "use_hipprof_profiler": True,
+        }
+
+    def test_injects_wrapper_and_preserves_custom_python(self):
+        result = configure_hipprof_env(
+            self.runner,
+            self.model,
+            {"PYTHON_EXEC": "/venv/bin/python"},
+        )
+
+        self.assertEqual(
+            result["PYTHON_EXEC"],
+            "tools/profiling/hipprof_python_wrapper.sh",
+        )
+        self.assertEqual(result["HIPPROF_REAL_PYTHON"], "/venv/bin/python")
+        self.assertEqual(result["HIPPROF_BIN_PATH"], "/opt/dtk/bin/hipprof")
+        self.assertEqual(result["HIPPROF_OUTPUT_DIR"], "/tmp/hipprof")
+
+    def test_accepts_full_executable_path(self):
+        self.runner["hipprof_bin_path"] = "/opt/dtk/bin/hipprof"
+
+        result = configure_hipprof_env(self.runner, self.model, {})
+
+        self.assertEqual(result["HIPPROF_BIN_PATH"], "/opt/dtk/bin/hipprof")
+
+    def test_disabled_configuration_is_unchanged(self):
+        result = configure_hipprof_env({}, {}, {"TOKEN": "value"})
+
+        self.assertEqual(result, {"TOKEN": "value"})
+
+    def test_launcher_only_profiling_arguments_are_removed(self):
+        runner_args = {
+            "nsys_bin_path": "/opt/nsight/bin/nsys",
+            "nsys_rep_file_path": "/tmp/nsys",
+            "hipprof_bin_path": "/opt/dtk/bin/hipprof",
+            "hipprof_output_dir": "/tmp/hipprof",
+            "nproc_per_node": 8,
+        }
+
+        remove_launcher_profiling_args(runner_args)
+
+        self.assertEqual(runner_args, {"nproc_per_node": 8})
+
+    def test_rejects_invalid_configurations(self):
+        cases = [
+            ({"hipprof_bin_path": "/opt/dtk/bin"}, self.model, "requires both"),
+            (
+                self.runner,
+                {"profile": True, "use_hipprof_profiler": False},
+                "use_hipprof_profiler: true",
+            ),
+            (
+                self.runner,
+                {"profile": False, "use_hipprof_profiler": True},
+                "profile: true",
+            ),
+            (
+                {
+                    **self.runner,
+                    "nsys_bin_path": "/opt/nsight/bin/nsys",
+                    "nsys_rep_file_path": "/tmp/nsys",
+                },
+                self.model,
+                "cannot be enabled together",
+            ),
+        ]
+        for runner, model, message in cases:
+            with (
+                self.subTest(message=message),
+                self.assertRaisesRegex(ValueError, message),
+            ):
+                configure_hipprof_env(runner, model, {})
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit_tests/runner/test_launcher_ssh_perf.py

@@ -37,3 +37,36 @@ def test_get_runner_cmd_train_strips_perf_monitor_runner_keys():
     assert "--perf_console_output" not in cmd
     assert "--log_dir" in cmd
     assert "--rdzv_endpoint" in cmd
+
+
+def test_get_runner_cmd_train_strips_hipprof_runner_keys():
+    config = OmegaConf.create(
+        {
+            "experiment": {
+                "runner": {
+                    "backend": "torchrun",
+                    "nnodes": 1,
+                    "nproc_per_node": 8,
+                    "rdzv_backend": "static",
+                    "hipprof_bin_path": "/opt/dtk/bin/hipprof",
+                    "hipprof_output_dir": "/tmp/hipprof",
+                }
+            },
+            "train": {
+                "system": {
+                    "logging": {
+                        "details_dir": "/tmp/details",
+                    }
+                },
+                "model": {
+                    "profile": True,
+                    "use_hipprof_profiler": True,
+                },
+            },
+        }
+    )
+
+    cmd = _get_runner_cmd_train("localhost", "127.0.0.1", 29500, 1, 0, 8, config)
+
+    assert "--hipprof_bin_path" not in cmd
+    assert "--hipprof_output_dir" not in cmd