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..d1b03786d0a 100644 --- a/tests/utils/test_torch_profile.py +++ b/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). diff --git a/verl/utils/profiler/config.py b/verl/utils/profiler/config.py index 569bb635c21..19bd3999ffa 100644 --- a/verl/utils/profiler/config.py +++ b/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}" diff --git a/verl/utils/profiler/torch_profile.py b/verl/utils/profiler/torch_profile.py index 0f8eafc7c79..f54d2922238 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_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,