[hardware, profile] fix: improve existing pytorch profiling to support other pytorch supported device

docs/perf/torch_profiling.md

@@ -33,7 +33,11 @@ You can customize the PyTorch Profiler behavior using the following fields under
         verl's per-stage markers are CPU-side events, so a device-only trace would be bare kernels
         with no way to tell which stage they belong to. Listing it is therefore redundant, and the
         rest of `contents` is honored as written.
-    *   **`cuda`**: Profile CUDA activities.
+    *   **`cuda`**: Profile CUDA activities. This is the device keyword on NVIDIA and on AMD
+        ROCm (PyTorch exposes HIP as `torch.cuda`). On another accelerator, use that platform's
+        device type instead — `xpu`, `hpu`, `mtia` — which is the one `torch.profiler` activity
+        that session records. Naming a device that is not the running platform's records no
+        device activity and logs a warning.
     *   **`memory`**: Track tensor memory allocation/free.
     *   **`shapes`**: Record shapes of operator inputs.
     *   **`stack`**: Record source code file and line number.

tests/utils/test_torch_profile.py

@@ -12,16 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import tempfile
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest import mock
+from unittest.mock import MagicMock, PropertyMock, patch
 
 import torch
 from omegaconf import OmegaConf
 
 from verl.utils.config import omega_conf_to_dataclass
 from verl.utils.profiler.config import (
+    DEVICE_ACTIVITIES,
+    DEVICE_CONTENTS,
     ProfilerConfig,
     TorchProfilerScheduleConfig,
     TorchProfilerToolConfig,
@@ -34,6 +38,25 @@
 )
 
 
+def _cuda_platform():
+    """A CUDA platform, for tests asserting the CUDA activity.
+
+    The device activity now comes from the active platform, so these tests are about a CUDA
+    platform specifically; patching it in keeps them passing on non-CUDA hardware.
+    """
+    platform = MagicMock()
+    type(platform).device_name = PropertyMock(return_value="cuda")
+    return platform
+
+
+# Some device torch can profile that is not CUDA, taken from whatever this build exposes so that no
+# single accelerator is baked into the tests.
+_OTHER_DEVICE = next(iter(sorted(device for device in DEVICE_ACTIVITIES if device != "cuda")), None)
+_requires_other_device = unittest.skipIf(
+    _OTHER_DEVICE is None, "this torch build exposes no device activity other than CUDA"
+)
+
+
 class TestTorchProfile(unittest.TestCase):
     def setUp(self):
         # Reset process-global Profiler class state so tests don't leak into each other.
@@ -44,6 +67,7 @@ def tearDown(self):
         Profiler._define_count = 0
         Profiler._active_prof = None
 
+    @patch("verl.utils.profiler.torch_profile.get_platform", new=_cuda_platform)
     @patch("torch.profiler.profile")
     def test_get_torch_profiler(self, mock_profile):
         # Test wrapper function
@@ -626,6 +650,7 @@ class TestCpuActivityAlwaysCollected(unittest.TestCase):
     built, so no configuration can produce a trace of unattributable kernels.
     """
 
+    @patch("verl.utils.profiler.torch_profile.get_platform", new=_cuda_platform)
     @patch("torch.profiler.profile")
     def test_cpu_activity_added_for_device_only_contents(self, mock_profile):
         for contents in (["cuda"], ["cuda", "memory"], []):
@@ -646,6 +671,122 @@ def test_contents_are_not_rewritten(self, mock_profile):
         self.assertEqual(contents, ["cuda"])
 
 
+class TestDeviceActivityFromPlatform(unittest.TestCase):
+    """The device activity is resolved from the active platform's ``device_name``."""
+
+    def _platform(self, device_name):
+        platform = MagicMock()
+        type(platform).device_name = PropertyMock(return_value=device_name)
+        return platform
+
+    def _activities(self, mock_profile):
+        return mock_profile.call_args[1]["activities"]
+
+    @_requires_other_device
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_non_cuda_device_recorded(self, mock_profile, mock_get_platform):
+        mock_get_platform.return_value = self._platform(_OTHER_DEVICE)
+        for contents in ([], [_OTHER_DEVICE], [_OTHER_DEVICE, "memory"]):
+            with self.subTest(contents=contents):
+                mock_profile.reset_mock()
+                get_torch_profiler(contents=list(contents), save_path="/tmp/test", rank=0)
+                activities = self._activities(mock_profile)
+                self.assertIn(DEVICE_ACTIVITIES[_OTHER_DEVICE], activities)
+                # torch checks CUDA before every other device: leaving it in a list it does not
+                # belong to disables device profiling altogether.
+                self.assertNotIn(torch.profiler.ProfilerActivity.CUDA, activities)
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_cuda_platform_unchanged(self, mock_profile, mock_get_platform):
+        # Regression guard for every in-tree CUDA/ROCm user: same activities as before.
+        mock_get_platform.return_value = self._platform("cuda")
+        for contents in ([], ["cuda"], ["cpu", "cuda", "stack"]):
+            with self.subTest(contents=contents):
+                mock_profile.reset_mock()
+                get_torch_profiler(contents=list(contents), save_path="/tmp/test", rank=0)
+                self.assertIn(torch.profiler.ProfilerActivity.CUDA, self._activities(mock_profile))
+
+    @_requires_other_device
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_other_platforms_device_keyword_warns(self, mock_profile, mock_get_platform):
+        # A CUDA config run on another accelerator records no device activity, so say so.
+        mock_get_platform.return_value = self._platform(_OTHER_DEVICE)
+        with self.assertLogs("verl.utils.profiler.torch_profile", level="WARNING") as log_ctx:
+            get_torch_profiler(contents=["cuda"], save_path="/tmp/test", rank=0)
+        self.assertEqual(self._activities(mock_profile), [torch.profiler.ProfilerActivity.CPU])
+        self.assertTrue(any(f"'{_OTHER_DEVICE}'" in record.getMessage() for record in log_ctx.records))
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_device_without_profiler_activity_warns(self, mock_profile, mock_get_platform):
+        # torch has no ProfilerActivity.NPU: same CPU-only trace as before, but the warning
+        # explaining it is now verl's instead of torch's misleading CUDA one.
+        mock_get_platform.return_value = self._platform("npu")
+        with self.assertLogs("verl.utils.profiler.torch_profile", level="WARNING"):
+            get_torch_profiler(contents=[], save_path="/tmp/test", rank=0)
+        self.assertEqual(self._activities(mock_profile), [torch.profiler.ProfilerActivity.CPU])
+
+    @_requires_other_device
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_non_device_contents_select_no_device(self, mock_profile, mock_get_platform):
+        # `contents` naming only profiler options is an explicit CPU-only request, not a mistake.
+        mock_get_platform.return_value = self._platform(_OTHER_DEVICE)
+        logger = logging.getLogger("verl.utils.profiler.torch_profile")
+        with mock.patch.object(logger, "warning") as mock_warning:
+            get_torch_profiler(contents=["stack"], save_path="/tmp/test", rank=0)
+        self.assertEqual(self._activities(mock_profile), [torch.profiler.ProfilerActivity.CPU])
+        mock_warning.assert_not_called()
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_every_device_keyword_resolves_its_activity(self, mock_profile, mock_get_platform):
+        # Holds for whatever members a torch build exposes, however they are capitalized.
+        for device, activity in DEVICE_ACTIVITIES.items():
+            with self.subTest(device=device):
+                mock_profile.reset_mock()
+                mock_get_platform.return_value = self._platform(device)
+                get_torch_profiler(contents=[], save_path="/tmp/test", rank=0)
+                self.assertIn(activity, self._activities(mock_profile))
+
+    @patch("verl.utils.profiler.torch_profile.get_platform")
+    @patch("torch.profiler.profile")
+    def test_cpu_device_not_added_twice(self, mock_profile, mock_get_platform):
+        mock_get_platform.return_value = self._platform("cpu")
+        get_torch_profiler(contents=[], save_path="/tmp/test", rank=0)
+        self.assertEqual(self._activities(mock_profile), [torch.profiler.ProfilerActivity.CPU])
+
+
+class TestDeviceContentsConfigValidation(unittest.TestCase):
+    """``contents`` accepts a device keyword for any backend torch can profile.
+
+    The keywords come from torch rather than a hardcoded list, so validation needs no platform
+    detection on the (often CPU-only) driver.
+    """
+
+    def test_every_torch_device_keyword_accepted(self):
+        for device in DEVICE_CONTENTS:
+            with self.subTest(device=device):
+                TorchProfilerToolConfig(contents=[device], discrete=False)
+
+    def test_profiler_options_still_accepted(self):
+        TorchProfilerToolConfig(contents=["cpu", "cuda", "memory", "shapes", "stack"], discrete=False)
+
+    def test_unknown_keyword_rejected(self):
+        with self.assertRaises(AssertionError):
+            TorchProfilerToolConfig(contents=["not_a_device"], discrete=False)
+
+    def test_private_use1_keyword_rejected(self):
+        # No platform reports "privateuse1" as its device_name, so accepting the keyword would
+        # validate a value that is then silently never recorded.
+        self.assertNotIn("privateuse1", DEVICE_CONTENTS)
+        with self.assertRaises(AssertionError):
+            TorchProfilerToolConfig(contents=["privateuse1"], discrete=False)
+
+
 def _role_profiler_omegaconf(tool="torch", enable=True, discrete=False, contents=("cpu", "cuda")):
     """Mimic a per-role ``profiler`` OmegaConf sub-tree (identical across ref/ref.yaml and
     critic/critic.yaml).

verl/utils/profiler/config.py

@@ -21,10 +21,24 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
+import torch
 from omegaconf import MISSING
 
 from verl.base_config import BaseConfig
 
+# Device type -> the `ProfilerActivity` recording it. torch names one member per backend it can
+# profile after that device type, so `get_torch_profiler` resolves the active platform's
+# `device_name` here. `CPU` is excluded (always collected, never resolved from a platform) and so is
+# `PrivateUse1` (torch's out-of-tree slot, not a device type any platform reports).
+DEVICE_ACTIVITIES = {
+    name.lower(): member
+    for name, member in torch.profiler.ProfilerActivity.__members__.items()
+    if name not in ("CPU", "PrivateUse1")
+}
+
+# `contents` keywords selecting a device activity rather than a profiler option.
+DEVICE_CONTENTS = frozenset(DEVICE_ACTIVITIES) | {"cpu"}
+
 
 @dataclass
 class NsightToolConfig(BaseConfig):
@@ -108,7 +122,8 @@ class TorchProfilerToolConfig(BaseConfig):
     (see :class:`TorchProfilerScheduleConfig`).
     """
 
-    # options: cuda, cpu, memory, shapes, stack. Empty means collect everything.
+    # options: the active platform's device type, cpu, memory, shapes, stack.
+    # Empty means collect everything.
     # CPU activity is collected either way (the per-stage record_function markers are CPU-side
     # events), so listing "cpu" here is redundant; the other options are honored as written.
     contents: list[str] = field(default_factory=list)
@@ -126,7 +141,7 @@ class TorchProfilerToolConfig(BaseConfig):
 
     def __post_init__(self) -> None:
         """config validation logics go here"""
-        __support_contents = ["cuda", "cpu", "memory", "shapes", "stack"]
+        __support_contents = sorted(DEVICE_CONTENTS | {"memory", "shapes", "stack"})
         for content in self.contents:
             assert content in __support_contents, (
                 f"Profiler contents only supports {__support_contents}, but gets {content}"

verl/utils/profiler/torch_profile.py

@@ -13,16 +13,21 @@
 # limitations under the License.
 
 import functools
+import logging
 import os
 import re
 from datetime import datetime, timezone
 from typing import Callable, Optional
 
 import torch
 
-from .config import ProfilerConfig, TorchProfilerToolConfig
+from verl.plugin.platform import get_platform
+
+from .config import DEVICE_ACTIVITIES, DEVICE_CONTENTS, ProfilerConfig, TorchProfilerToolConfig
 from .profile import DistProfiler
 
+logger = logging.getLogger(__name__)
+
 
 def get_dist_topology() -> dict:
     """Best-effort snapshot of the current process's distributed topology.
@@ -122,10 +127,11 @@ def get_torch_profiler(
     """Build a ``torch.profiler.profile`` instance.
 
     Args:
-        contents: Selects the other ``torch.profiler.profile`` arguments -- ``cuda`` maps to
-            ``activities``, ``shapes`` to ``record_shapes``, ``memory`` to ``profile_memory`` and
-            ``stack`` to ``with_stack``. CPU activity is always on, since verl's per-stage
-            ``record_function`` markers are CPU-side events.
+        contents: Selects the other ``torch.profiler.profile`` arguments -- the active platform's
+            device type (``cuda`` on NVIDIA and ROCm) maps to ``activities``, ``shapes`` to
+            ``record_shapes``, ``memory`` to ``profile_memory`` and ``stack`` to ``with_stack``.
+            CPU activity is always on, since verl's per-stage ``record_function`` markers are
+            CPU-side events.
         save_path: Directory to write chrome traces to.
         role: Optional logical scope name (e.g. ``train`` for a worker's whole-step window, or
             a stage name in discrete mode), embedded in the filename.
@@ -196,8 +202,29 @@ def _trace_handler(prof):
     # record_function, and those markers -- like operator names -- are CPU-side events, so a
     # device-only trace would be bare kernels that cannot be attributed to any stage.
     activities = [torch.profiler.ProfilerActivity.CPU]
-    if not contents or "cuda" in contents:
-        activities.append(torch.profiler.ProfilerActivity.CUDA)
+    # The platform's own device type is both the activity to record and the `contents` keyword that
+    # asks for it, so no per-backend branch is needed: NVIDIA/ROCm resolve "cuda" exactly as before,
+    # and a device torch cannot profile resolves to nothing.
+    device = get_platform().device_name
+    device_activity = DEVICE_ACTIVITIES.get(device)
+    device_requested = not contents or device in contents
+    if device_requested and device_activity is not None:
+        activities.append(device_activity)
+    elif device_requested and device != "cpu":
+        logger.warning(
+            "torch.profiler has no activity for device '%s', so the trace holds CPU events only. "
+            "If this platform has its own profiler, select it with global_profiler.tool instead.",
+            device,
+        )
+    elif ignored := sorted((contents & DEVICE_CONTENTS) - {device, "cpu"}):
+        # A keyword valid for *a* backend but not this one, e.g. a CUDA config run elsewhere.
+        logger.warning(
+            "profiler contents selects device activity %s, but this platform's device is '%s': "
+            "no device activity is recorded. Use '%s', or leave contents empty.",
+            ignored,
+            device,
+            device,
+        )
 
     profile_kwargs = dict(
         activities=activities,