diff --git a/tests/plugin/test_platform_abstraction.py b/tests/plugin/test_platform_abstraction.py index d4a88618438..046f9a17e51 100644 --- a/tests/plugin/test_platform_abstraction.py +++ b/tests/plugin/test_platform_abstraction.py @@ -208,5 +208,14 @@ def test_unregistered_platform_raises(self): _create_platform("nonexistent_platform") +class TestTorchProfilerHooks: + """Test the default (no-op) behavior of the torch.profiler plugin hooks.""" + + def test_defaults_are_none(self): + platform = _make_mock_platform("mock_torch_profiler")() + assert platform.torch_profiler_activity() is None + assert platform.torch_profiler_content_name() is None + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/utils/test_torch_profile.py b/tests/utils/test_torch_profile.py index e38bca9ed39..79e07943c27 100644 --- a/tests/utils/test_torch_profile.py +++ b/tests/utils/test_torch_profile.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import os import tempfile import unittest +from unittest import mock from unittest.mock import MagicMock, patch import torch @@ -646,6 +648,119 @@ def test_contents_are_not_rewritten(self, mock_profile): self.assertEqual(contents, ["cuda"]) +class TestPluginProfilerActivityHook(unittest.TestCase): + """A platform-supplied ``torch.profiler`` activity is added only when both + ``torch_profiler_activity()`` and ``torch_profiler_content_name()`` return non-``None``, and + (if ``contents`` is non-empty) the content name is one of the requested keywords. CUDA must + stay enabled when the user explicitly asks for it, even on a platform that also exposes its + own activity hook (regression guard for the plugin-device-availability vs. + plugin-device-requested mixup).""" + + # A vendor-neutral stand-in: the code under test never interprets this value, it only + # stores/compares it, so any sentinel distinct from the real ProfilerActivity members works. + _MOCK_ACTIVITY = object() + + def _mock_platform(self, activity, content_name): + platform = MagicMock() + platform.torch_profiler_activity.return_value = activity + platform.torch_profiler_content_name.return_value = content_name + return platform + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_added_for_empty_or_matching_contents(self, mock_profile, mock_get_platform): + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device") + for contents in ([], ["mock_device"]): + with self.subTest(contents=contents): + mock_profile.reset_mock() + get_torch_profiler(contents=list(contents), save_path="/tmp/test", rank=0) + activities = mock_profile.call_args[1]["activities"] + self.assertIn(self._MOCK_ACTIVITY, activities) + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_excluded_for_non_matching_contents(self, mock_profile, mock_get_platform): + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device") + get_torch_profiler(contents=["stack"], save_path="/tmp/test", rank=0) + activities = mock_profile.call_args[1]["activities"] + self.assertNotIn(self._MOCK_ACTIVITY, activities) + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_excluded_when_content_name_missing(self, mock_profile, mock_get_platform): + # Regression guard: an activity hook without a paired content name must never be added, + # even for empty `contents`. + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, None) + get_torch_profiler(contents=[], save_path="/tmp/test", rank=0) + activities = mock_profile.call_args[1]["activities"] + self.assertNotIn(self._MOCK_ACTIVITY, activities) + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_excluded_when_activity_missing(self, mock_profile, mock_get_platform): + mock_get_platform.return_value = self._mock_platform(None, "mock_device") + get_torch_profiler(contents=["mock_device"], save_path="/tmp/test", rank=0) + activities = mock_profile.call_args[1]["activities"] + self.assertEqual(activities, [torch.profiler.ProfilerActivity.CPU]) + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_explicit_cuda_wins_over_available_plugin_device(self, mock_profile, mock_get_platform): + # A plugin platform merely exposing an activity must not steal `contents=["cuda"]` away + # from CUDA: the plugin device has to be *requested*, not just *available*. + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device") + get_torch_profiler(contents=["cuda"], save_path="/tmp/test", rank=0) + activities = mock_profile.call_args[1]["activities"] + self.assertIn(torch.profiler.ProfilerActivity.CUDA, activities) + self.assertNotIn(self._MOCK_ACTIVITY, activities) + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_warns_when_cuda_is_requested_but_superseded(self, mock_profile, mock_get_platform): + # contents=["cuda", ] is self-contradictory: only one activity is ever + # recorded, and there is no active CUDA platform here, so "cuda" is silently dead unless + # we say so. + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device") + with self.assertLogs("verl.utils.profiler.torch_profile", level="WARNING") as log_ctx: + get_torch_profiler(contents=["cuda", "mock_device"], save_path="/tmp/test", rank=0) + self.assertTrue(any("ignored" in record.getMessage() for record in log_ctx.records)) + + @patch("verl.utils.profiler.torch_profile.get_platform") + @patch("torch.profiler.profile") + def test_no_warning_without_explicit_cuda(self, mock_profile, mock_get_platform): + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device") + logger = logging.getLogger("verl.utils.profiler.torch_profile") + with mock.patch.object(logger, "warning") as mock_warning: + get_torch_profiler(contents=["mock_device"], save_path="/tmp/test", rank=0) + mock_warning.assert_not_called() + + +class TestPluginContentNameConfigValidation(unittest.TestCase): + """``TorchProfilerToolConfig.__post_init__`` must only whitelist a platform's + ``torch_profiler_content_name()`` keyword when ``torch_profiler_activity()`` is also + non-``None`` -- otherwise a config naming that keyword passes validation here but + ``get_torch_profiler()`` silently drops it (see ``TestPluginProfilerActivityHook``).""" + + _MOCK_ACTIVITY = object() + + def _mock_platform(self, activity, content_name): + platform = MagicMock() + platform.torch_profiler_activity.return_value = activity + platform.torch_profiler_content_name.return_value = content_name + return platform + + @patch("verl.utils.profiler.config.get_platform") + def test_content_name_allowed_when_activity_present(self, mock_get_platform): + mock_get_platform.return_value = self._mock_platform(self._MOCK_ACTIVITY, "mock_device") + TorchProfilerToolConfig(contents=["mock_device"], discrete=False) + + @patch("verl.utils.profiler.config.get_platform") + def test_content_name_rejected_when_activity_missing(self, mock_get_platform): + mock_get_platform.return_value = self._mock_platform(None, "mock_device") + with self.assertRaises(AssertionError): + TorchProfilerToolConfig(contents=["mock_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/plugin/platform/README.md b/verl/plugin/platform/README.md index 38fc89ae8e3..a65ae6c9ff0 100644 --- a/verl/plugin/platform/README.md +++ b/verl/plugin/platform/README.md @@ -230,6 +230,8 @@ VERL_USE_EXTERNAL_PLUGINS=pkg1,pkg2 # only load named entry_points | **Profiling** | `nvtx_range(msg)` | Context manager for profiler ranges | | | `profiler_start()` | Start device profiler | | | `profiler_stop()` | Stop device profiler | +| | `torch_profiler_activity()` | This platform's `torch.profiler.ProfilerActivity`, or `None` to opt out | +| | `torch_profiler_content_name()` | `contents` keyword that requests it (e.g. `'xpu'`); must be paired with `torch_profiler_activity()` | | **Low-level** | `cudart()` | CUDA runtime API object or `None` | | | `check_smi_command(cmd)` | Run SMI command and check exit code | diff --git a/verl/plugin/platform/platform_base.py b/verl/plugin/platform/platform_base.py index 4ae0c1a098a..f0e42984119 100644 --- a/verl/plugin/platform/platform_base.py +++ b/verl/plugin/platform/platform_base.py @@ -14,7 +14,10 @@ import subprocess from contextlib import contextmanager from types import ModuleType -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional + +if TYPE_CHECKING: + import torch class PlatformBase(abc.ABC): @@ -189,6 +192,27 @@ def profiler_stop(self) -> None: """Stop the device profiler (no-op on unsupported platforms).""" ... + def torch_profiler_activity(self) -> Optional["torch.profiler.ProfilerActivity"]: + """Return this platform's ``torch.profiler`` device activity, if any. + + Lets a platform request its own device-activity be collected by + ``verl/utils/profiler/torch_profile.py::get_torch_profiler`` when the + matching name (see :meth:`torch_profiler_content_name`) appears in + ``profiler.tool_config.torch.contents``. Return ``None`` (default) if + this platform has no ``torch.profiler.ProfilerActivity`` member (the + ``cuda`` activity is handled separately and is unaffected by this hook). + """ + return None + + def torch_profiler_content_name(self) -> Optional[str]: + """Return the ``contents`` keyword that requests this platform's activity. + + E.g. ``"xpu"`` for Intel XPU. Paired with :meth:`torch_profiler_activity`; + both must be non-``None`` for the activity to be added. Return ``None`` + (default) to opt out. + """ + return None + # ------------------------------------------------------------------ # vllm integration # ------------------------------------------------------------------ diff --git a/verl/utils/profiler/config.py b/verl/utils/profiler/config.py index 569bb635c21..bfe3debc7f6 100644 --- a/verl/utils/profiler/config.py +++ b/verl/utils/profiler/config.py @@ -24,6 +24,7 @@ from omegaconf import MISSING from verl.base_config import BaseConfig +from verl.plugin.platform import get_platform @dataclass @@ -127,6 +128,10 @@ class TorchProfilerToolConfig(BaseConfig): def __post_init__(self) -> None: """config validation logics go here""" __support_contents = ["cuda", "cpu", "memory", "shapes", "stack"] + plugin_content = get_platform().torch_profiler_content_name() + plugin_activity = get_platform().torch_profiler_activity() + if plugin_content is not None and plugin_activity is not None: + __support_contents = [*__support_contents, plugin_content] 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..64168893c7b 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 verl.plugin.platform import get_platform + from .config import 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. @@ -125,7 +130,9 @@ def get_torch_profiler( 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. + ``record_function`` markers are CPU-side events. A plugin platform may also accept + its own content name here (see ``PlatformBase.torch_profiler_content_name``), mapped + to that platform's own ``torch.profiler.ProfilerActivity``. 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 +203,27 @@ 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: + plugin_activity = get_platform().torch_profiler_activity() + plugin_content = get_platform().torch_profiler_content_name() + has_plugin_device = plugin_activity is not None and plugin_content is not None + plugin_requested = has_plugin_device and (not contents or plugin_content in contents) + # torch.profiler records a single device type and checks CUDA before every other device, + # so CUDA must stay out of `activities` whenever the plugin device was actually requested. + if not plugin_requested and (not contents or "cuda" in contents): activities.append(torch.profiler.ProfilerActivity.CUDA) + if plugin_requested: + activities.append(plugin_activity) + if "cuda" in contents: + # Only one device activity is ever recorded, and there is no active CUDA + # platform here (CUDA never implements these hooks) -- "cuda" is a dead + # keyword in this contents list, not a second activity. + logger.warning( + "profiler contents=%s requested both 'cuda' and '%s': only '%s' " + "(this platform's activity) is recorded, 'cuda' is ignored.", + sorted(contents), + plugin_content, + plugin_content, + ) profile_kwargs = dict( activities=activities,