From ee10d62c89bf3b2ea8a9b07c5b6fc8807cdb9101 Mon Sep 17 00:00:00 2001 From: kahlun Date: Thu, 8 Oct 2026 08:57:59 +0000 Subject: [PATCH 1/3] [profile] refactor: resolve torch.profiler device activity from the active platform ### What does this PR do? `get_torch_profiler` knew exactly one device activity, `ProfilerActivity.CUDA`, and `TorchProfilerToolConfig` accepted exactly one device keyword, `"cuda"`. On an accelerator that is not CUDA there was then no way to record device kernels: the default `contents=[]` asks for CUDA, torch logs "CUDA is not available, disabling CUDA profiling" and the trace holds CPU ops only, while naming the real device in `contents` fails config validation. torch names each `ProfilerActivity` member after the device type it profiles, and a single profiler session records exactly one of them. So the active platform's `device_name` is both the activity to collect and the `contents` keyword that requests it. Resolve it that way instead of hardcoding CUDA, and take the accepted device keywords from `torch.profiler.ProfilerActivity` rather than from a verl-maintained list. No new `PlatformBase` API: `device_name` is already an abstract property. ### No change for existing platforms `PlatformCUDA.device_name` is `"cuda"`, and `PlatformROCm`, `PlatformMetaX` and `PlatformIluvatar` inherit it, so all four keep building the same `activities` list. Platforms with no matching member (`npu`, `mlu`, `musa`, `gcu`, `tpu`, `supa`) still get a CPU-only trace, exactly as they do today when torch disables the unavailable CUDA activity -- the warning explaining it is now verl's, and it names the device. ### Tests `test_get_torch_profiler` and `test_cpu_activity_added_for_device_only_contents` assert the CUDA activity, which until now was hardcoded and so held on any machine. They now patch in a CUDA platform, stating the assumption they always had -- without it they fail when the suite is run on non-CUDA hardware. Verified by running the file with `VERL_PLATFORM` forced to a non-CUDA platform. --- docs/perf/torch_profiling.md | 6 +- tests/utils/test_torch_profile.py | 122 ++++++++++++++++++++++++++- verl/utils/profiler/config.py | 14 ++- verl/utils/profiler/torch_profile.py | 46 ++++++++-- 4 files changed, 177 insertions(+), 11 deletions(-) diff --git a/docs/perf/torch_profiling.md b/docs/perf/torch_profiling.md index 04d0903ba3b..31083e0131a 100644 --- a/docs/perf/torch_profiling.md +++ b/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. diff --git a/tests/utils/test_torch_profile.py b/tests/utils/test_torch_profile.py index e38bca9ed39..404914d05dd 100644 --- a/tests/utils/test_torch_profile.py +++ b/tests/utils/test_torch_profile.py @@ -12,10 +12,12 @@ # 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 @@ -34,6 +36,19 @@ ) +def _cuda_platform(): + """A platform whose device is CUDA, for tests that assert the CUDA activity. + + ``get_torch_profiler`` derives the device activity from the active platform, so a test that + names ``cuda`` in ``contents`` and expects ``ProfilerActivity.CUDA`` is asserting about a + CUDA platform specifically. Patching it in keeps those tests passing when the suite is run + on non-CUDA hardware, where ``get_platform()`` is that vendor's platform. + """ + platform = MagicMock() + type(platform).device_name = PropertyMock(return_value="cuda") + return platform + + class TestTorchProfile(unittest.TestCase): def setUp(self): # Reset process-global Profiler class state so tests don't leak into each other. @@ -44,6 +59,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 +642,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 +663,109 @@ 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``. + + torch names each ``ProfilerActivity`` member after the device type it profiles, so a backend + that has a member needs no verl change to be traced, and the NVIDIA/ROCm behaviour falls out + of the same rule (``PlatformCUDA.device_name`` is ``"cuda"``, and ``PlatformROCm`` inherits it). + """ + + 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"] + + @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("xpu") + for contents in ([], ["xpu"], ["xpu", "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(torch.profiler.ProfilerActivity.XPU, activities) + # Only one device activity is ever recorded, and CUDA is checked before every + # other device: leaving it in would disable device profiling altogether here. + 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)) + + @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 config written for CUDA, run on another accelerator: "cuda" is a valid keyword for + # *a* backend but not for this one, so no device activity is recorded. Say so. + mock_get_platform.return_value = self._platform("xpu") + 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("'xpu'" 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, so there is nothing to request. Previously verl + # asked for CUDA here and torch logged "CUDA is not available, disabling CUDA + # profiling"; the trace content is the same, the warning is now ours. + 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]) + + @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("xpu") + 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_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.profiler.ProfilerActivity`` rather than a hardcoded list, so + validation does not have to know which platform the run will land on -- which keeps + ``TorchProfilerToolConfig`` free of platform detection on the (often CPU-only) driver. + """ + + def test_every_torch_device_keyword_accepted(self): + for name in torch.profiler.ProfilerActivity.__members__: + with self.subTest(device=name): + TorchProfilerToolConfig(contents=[name.lower()], 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 _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). diff --git a/verl/utils/profiler/config.py b/verl/utils/profiler/config.py index 569bb635c21..9025362ad93 100644 --- a/verl/utils/profiler/config.py +++ b/verl/utils/profiler/config.py @@ -21,10 +21,19 @@ from dataclasses import dataclass, field from typing import Any, Optional +import torch from omegaconf import MISSING from verl.base_config import BaseConfig +# `contents` keywords that select a device activity rather than a profiler option. torch exposes +# one `ProfilerActivity` member per backend it can profile, named as the uppercased device type +# ("cuda", "xpu", "hpu", "mtia", ...), and a single profiler session records exactly one of them. +# Taking the keywords from torch keeps this list from being a verl-maintained copy that silently +# omits every backend but CUDA; `get_torch_profiler` resolves the active platform's `device_name` +# against the same set to pick which one is actually recorded. +DEVICE_CONTENTS = frozenset(name.lower() for name in torch.profiler.ProfilerActivity.__members__) + @dataclass class NsightToolConfig(BaseConfig): @@ -108,7 +117,8 @@ class TorchProfilerToolConfig(BaseConfig): (see :class:`TorchProfilerScheduleConfig`). """ - # options: cuda, cpu, memory, shapes, stack. Empty means collect everything. + # options: the active platform's device type (cuda, xpu, ...), 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 +136,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}" diff --git a/verl/utils/profiler/torch_profile.py b/verl/utils/profiler/torch_profile.py index 0f8eafc7c79..e1428dea430 100644 --- a/verl/utils/profiler/torch_profile.py +++ b/verl/utils/profiler/torch_profile.py @@ -13,6 +13,7 @@ # limitations under the License. import functools +import logging import os import re from datetime import datetime, timezone @@ -20,9 +21,13 @@ import torch -from .config import ProfilerConfig, TorchProfilerToolConfig +from verl.plugin.platform import get_platform + +from .config import 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, ``xpu``, ... ) 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,34 @@ 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 device activity to record is the active platform's own device type: torch names each + # `ProfilerActivity` member after the device type it profiles, so the platform's `device_name` + # is both the activity to collect and the `contents` keyword that asks for it. No per-backend + # branch is needed, and NVIDIA/ROCm keep the previous behaviour -- `PlatformCUDA.device_name` + # is "cuda" and `PlatformROCm` inherits it, since torch exposes HIP as `torch.cuda`. + device = get_platform().device_name + # "cpu" is not a second activity here: CPU is already in the list above. + device_activity = getattr(torch.profiler.ProfilerActivity, device.upper(), None) if device != "cpu" else None + 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 device keyword that is valid for *a* backend but not for this one (e.g. a config + # written for CUDA, run on another accelerator). Only one device activity is ever + # recorded, so the trace would hold CPU events alone with nothing explaining why. + 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, From c4661547aeba180d338b915876d041abf4de268c Mon Sep 17 00:00:00 2001 From: kahlun Date: Thu, 8 Oct 2026 12:00:24 +0000 Subject: [PATCH 2/3] [profile] fix: resolve the device activity through a device-type map `getattr(torch.profiler.ProfilerActivity, device.upper(), None)` depends on torch capitalizing every member as the uppercased device type, which `PrivateUse1` already breaks, and `DEVICE_CONTENTS` lowercased every member name -- so `contents: ["privateuse1"]` passed config validation and was then never recorded, since no platform reports "privateuse1" as its `device_name`. Resolve the activity through an explicit lowercased device-type -> activity map instead, and build the accepted `contents` keywords from it. `CPU` and `PrivateUse1` are left out of the map: neither is a device type a platform reports, CPU activity is collected in every trace regardless of `contents`, and "cpu" stays an accepted keyword. An out-of-tree backend whose torch activity is not named after its device type can add its own entry to the map, rather than verl carrying a hook for it. Tests: every entry in the map resolves to its own activity (so the lookup holds for whatever members a torch build exposes), and "privateuse1" is rejected by config validation. --- tests/utils/test_torch_profile.py | 28 +++++++++++++++++++++++++--- verl/utils/profiler/config.py | 27 ++++++++++++++++++++------- verl/utils/profiler/torch_profile.py | 7 ++++--- 3 files changed, 49 insertions(+), 13 deletions(-) diff --git a/tests/utils/test_torch_profile.py b/tests/utils/test_torch_profile.py index 404914d05dd..31b86678fae 100644 --- a/tests/utils/test_torch_profile.py +++ b/tests/utils/test_torch_profile.py @@ -24,6 +24,8 @@ from verl.utils.config import omega_conf_to_dataclass from verl.utils.profiler.config import ( + DEVICE_ACTIVITIES, + DEVICE_CONTENTS, ProfilerConfig, TorchProfilerScheduleConfig, TorchProfilerToolConfig, @@ -737,6 +739,18 @@ def test_non_device_contents_select_no_device(self, mock_profile, mock_get_platf 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): + # Every device type torch can profile resolves to its own activity, whatever this torch + # build happens to expose and however it capitalizes the member name. + 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): @@ -754,9 +768,9 @@ class TestDeviceContentsConfigValidation(unittest.TestCase): """ def test_every_torch_device_keyword_accepted(self): - for name in torch.profiler.ProfilerActivity.__members__: - with self.subTest(device=name): - TorchProfilerToolConfig(contents=[name.lower()], discrete=False) + 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) @@ -765,6 +779,14 @@ def test_unknown_keyword_rejected(self): with self.assertRaises(AssertionError): TorchProfilerToolConfig(contents=["not_a_device"], discrete=False) + def test_private_use1_keyword_rejected(self): + # `PrivateUse1` is torch's slot for an out-of-tree backend, not a device type any platform + # reports as its `device_name`, so no run could ever select it. 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 diff --git a/verl/utils/profiler/config.py b/verl/utils/profiler/config.py index 9025362ad93..dc62d7d887b 100644 --- a/verl/utils/profiler/config.py +++ b/verl/utils/profiler/config.py @@ -26,13 +26,26 @@ from verl.base_config import BaseConfig -# `contents` keywords that select a device activity rather than a profiler option. torch exposes -# one `ProfilerActivity` member per backend it can profile, named as the uppercased device type -# ("cuda", "xpu", "hpu", "mtia", ...), and a single profiler session records exactly one of them. -# Taking the keywords from torch keeps this list from being a verl-maintained copy that silently -# omits every backend but CUDA; `get_torch_profiler` resolves the active platform's `device_name` -# against the same set to pick which one is actually recorded. -DEVICE_CONTENTS = frozenset(name.lower() for name in torch.profiler.ProfilerActivity.__members__) +# Device type -> the `ProfilerActivity` that records its kernels. torch exposes one member per +# backend it can profile, named after the device type ("cuda", "xpu", "hpu", "mtia", ...), and a +# single profiler session records exactly one of them. Keying by the lowercased name keeps the +# lookup independent of how torch capitalizes a member, and taking the set from torch keeps it from +# being a verl-maintained copy that silently omits every backend but CUDA; `get_torch_profiler` +# resolves the active platform's `device_name` here to pick the activity it records. +# +# Two members are not device types and so are never resolved from a platform: `CPU`, which is +# collected in every trace regardless of `contents`, and `PrivateUse1`, which is torch's slot for an +# out-of-tree backend rather than a device any platform reports as its `device_name`. +DEVICE_ACTIVITIES = { + name.lower(): member + for name, member in torch.profiler.ProfilerActivity.__members__.items() + if name not in ("CPU", "PrivateUse1") +} + +# `contents` keywords that select a device activity rather than a profiler option. "cpu" is accepted +# for backwards compatibility and because it reads naturally, even though CPU activity is collected +# either way. +DEVICE_CONTENTS = frozenset(DEVICE_ACTIVITIES) | {"cpu"} @dataclass diff --git a/verl/utils/profiler/torch_profile.py b/verl/utils/profiler/torch_profile.py index e1428dea430..9fbd22232ca 100644 --- a/verl/utils/profiler/torch_profile.py +++ b/verl/utils/profiler/torch_profile.py @@ -23,7 +23,7 @@ from verl.plugin.platform import get_platform -from .config import DEVICE_CONTENTS, ProfilerConfig, TorchProfilerToolConfig +from .config import DEVICE_ACTIVITIES, DEVICE_CONTENTS, ProfilerConfig, TorchProfilerToolConfig from .profile import DistProfiler logger = logging.getLogger(__name__) @@ -208,8 +208,9 @@ def _trace_handler(prof): # branch is needed, and NVIDIA/ROCm keep the previous behaviour -- `PlatformCUDA.device_name` # is "cuda" and `PlatformROCm` inherits it, since torch exposes HIP as `torch.cuda`. device = get_platform().device_name - # "cpu" is not a second activity here: CPU is already in the list above. - device_activity = getattr(torch.profiler.ProfilerActivity, device.upper(), None) if device != "cpu" else None + # DEVICE_ACTIVITIES holds device types only, so a CPU-only platform resolves to no activity: + # CPU is already in the list above and is never added twice. + 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) From 771d453cd2b95f7f4c562ff0043a83d386d52e3d Mon Sep 17 00:00:00 2001 From: kahlun Date: Thu, 8 Oct 2026 12:25:28 +0000 Subject: [PATCH 3/3] make the docstring and comment more device agnostic --- tests/utils/test_torch_profile.py | 61 ++++++++++++++-------------- verl/utils/profiler/config.py | 20 +++------ verl/utils/profiler/torch_profile.py | 22 ++++------ 3 files changed, 44 insertions(+), 59 deletions(-) diff --git a/tests/utils/test_torch_profile.py b/tests/utils/test_torch_profile.py index 31b86678fae..d1b03786d0a 100644 --- a/tests/utils/test_torch_profile.py +++ b/tests/utils/test_torch_profile.py @@ -39,18 +39,24 @@ def _cuda_platform(): - """A platform whose device is CUDA, for tests that assert the CUDA activity. + """A CUDA platform, for tests asserting the CUDA activity. - ``get_torch_profiler`` derives the device activity from the active platform, so a test that - names ``cuda`` in ``contents`` and expects ``ProfilerActivity.CUDA`` is asserting about a - CUDA platform specifically. Patching it in keeps those tests passing when the suite is run - on non-CUDA hardware, where ``get_platform()`` is that vendor's platform. + 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. @@ -666,12 +672,7 @@ def test_contents_are_not_rewritten(self, mock_profile): class TestDeviceActivityFromPlatform(unittest.TestCase): - """The device activity is resolved from the active platform's ``device_name``. - - torch names each ``ProfilerActivity`` member after the device type it profiles, so a backend - that has a member needs no verl change to be traced, and the NVIDIA/ROCm behaviour falls out - of the same rule (``PlatformCUDA.device_name`` is ``"cuda"``, and ``PlatformROCm`` inherits it). - """ + """The device activity is resolved from the active platform's ``device_name``.""" def _platform(self, device_name): platform = MagicMock() @@ -681,18 +682,19 @@ def _platform(self, device_name): 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("xpu") - for contents in ([], ["xpu"], ["xpu", "memory"]): + 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(torch.profiler.ProfilerActivity.XPU, activities) - # Only one device activity is ever recorded, and CUDA is checked before every - # other device: leaving it in would disable device profiling altogether here. + 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") @@ -706,33 +708,33 @@ def test_cuda_platform_unchanged(self, mock_profile, mock_get_platform): 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 config written for CUDA, run on another accelerator: "cuda" is a valid keyword for - # *a* backend but not for this one, so no device activity is recorded. Say so. - mock_get_platform.return_value = self._platform("xpu") + # 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("'xpu'" in record.getMessage() for record in log_ctx.records)) + 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, so there is nothing to request. Previously verl - # asked for CUDA here and torch logged "CUDA is not available, disabling CUDA - # profiling"; the trace content is the same, the warning is now ours. + # 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("xpu") + 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) @@ -742,8 +744,7 @@ def test_non_device_contents_select_no_device(self, mock_profile, mock_get_platf @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): - # Every device type torch can profile resolves to its own activity, whatever this torch - # build happens to expose and however it capitalizes the member name. + # 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() @@ -762,9 +763,8 @@ def test_cpu_device_not_added_twice(self, mock_profile, mock_get_platform): class TestDeviceContentsConfigValidation(unittest.TestCase): """``contents`` accepts a device keyword for any backend torch can profile. - The keywords come from ``torch.profiler.ProfilerActivity`` rather than a hardcoded list, so - validation does not have to know which platform the run will land on -- which keeps - ``TorchProfilerToolConfig`` free of platform detection on the (often CPU-only) driver. + 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): @@ -780,8 +780,7 @@ def test_unknown_keyword_rejected(self): TorchProfilerToolConfig(contents=["not_a_device"], discrete=False) def test_private_use1_keyword_rejected(self): - # `PrivateUse1` is torch's slot for an out-of-tree backend, not a device type any platform - # reports as its `device_name`, so no run could ever select it. Accepting the keyword would + # 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): diff --git a/verl/utils/profiler/config.py b/verl/utils/profiler/config.py index dc62d7d887b..19bd3999ffa 100644 --- a/verl/utils/profiler/config.py +++ b/verl/utils/profiler/config.py @@ -26,25 +26,17 @@ from verl.base_config import BaseConfig -# Device type -> the `ProfilerActivity` that records its kernels. torch exposes one member per -# backend it can profile, named after the device type ("cuda", "xpu", "hpu", "mtia", ...), and a -# single profiler session records exactly one of them. Keying by the lowercased name keeps the -# lookup independent of how torch capitalizes a member, and taking the set from torch keeps it from -# being a verl-maintained copy that silently omits every backend but CUDA; `get_torch_profiler` -# resolves the active platform's `device_name` here to pick the activity it records. -# -# Two members are not device types and so are never resolved from a platform: `CPU`, which is -# collected in every trace regardless of `contents`, and `PrivateUse1`, which is torch's slot for an -# out-of-tree backend rather than a device any platform reports as its `device_name`. +# 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 that select a device activity rather than a profiler option. "cpu" is accepted -# for backwards compatibility and because it reads naturally, even though CPU activity is collected -# either way. +# `contents` keywords selecting a device activity rather than a profiler option. DEVICE_CONTENTS = frozenset(DEVICE_ACTIVITIES) | {"cpu"} @@ -130,7 +122,7 @@ class TorchProfilerToolConfig(BaseConfig): (see :class:`TorchProfilerScheduleConfig`). """ - # options: the active platform's device type (cuda, xpu, ...), cpu, memory, shapes, stack. + # 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. diff --git a/verl/utils/profiler/torch_profile.py b/verl/utils/profiler/torch_profile.py index 9fbd22232ca..f54d2922238 100644 --- a/verl/utils/profiler/torch_profile.py +++ b/verl/utils/profiler/torch_profile.py @@ -128,10 +128,10 @@ def get_torch_profiler( Args: contents: Selects the other ``torch.profiler.profile`` arguments -- the active platform's - device type (``cuda`` on NVIDIA and ROCm, ``xpu``, ... ) 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. + 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. @@ -202,14 +202,10 @@ 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] - # The device activity to record is the active platform's own device type: torch names each - # `ProfilerActivity` member after the device type it profiles, so the platform's `device_name` - # is both the activity to collect and the `contents` keyword that asks for it. No per-backend - # branch is needed, and NVIDIA/ROCm keep the previous behaviour -- `PlatformCUDA.device_name` - # is "cuda" and `PlatformROCm` inherits it, since torch exposes HIP as `torch.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_ACTIVITIES holds device types only, so a CPU-only platform resolves to no activity: - # CPU is already in the list above and is never added twice. device_activity = DEVICE_ACTIVITIES.get(device) device_requested = not contents or device in contents if device_requested and device_activity is not None: @@ -221,9 +217,7 @@ def _trace_handler(prof): device, ) elif ignored := sorted((contents & DEVICE_CONTENTS) - {device, "cpu"}): - # A device keyword that is valid for *a* backend but not for this one (e.g. a config - # written for CUDA, run on another accelerator). Only one device activity is ever - # recorded, so the trace would hold CPU events alone with nothing explaining why. + # 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.",