diff --git a/docs/design/xpu-monkeypatch-experiment.md b/docs/design/xpu-monkeypatch-experiment.md new file mode 100644 index 00000000..fa351b46 --- /dev/null +++ b/docs/design/xpu-monkeypatch-experiment.md @@ -0,0 +1,122 @@ +# Experiment: how much of #7917 can live in the plugin as pure monkeypatch? + +## Context + +`verl-hardware-plugin#26` asked maintainers whether Intel XPU should extend +`PlatformBase` in `verl-core` (the TPU precedent) or stay entirely inside +this plugin via monkeypatch (the MLU precedent). The maintainer's answer +(2026-09-23): default to the plugin, keep `verl-project/verl#7917` (which +adds `is_reduce_avg_supported`/`attention_utils_module`/`profiler_markers`/ +`dist_profiler_cls` to `PlatformBase`) narrowed to only what truly can't be +done from here. + +This branch is the answer to "then how much of #7917 actually can't be done +from here" — implemented, not just argued. It does not touch +`feature/xpu-vtune-avg-attention` (the branch behind open PR #22, which +still uses the hook-based approach and depends on #7917 merging). This is a +separate, from-`main` branch so the two can be compared directly for size +and how much they depend on core changes ever landing. + +## Verdict per capability + +| Capability | In this branch? | Mechanism | Notes | +|---|---|---|---| +| `attention_utils_module` | Yes | Patch `verl.utils.attention_utils._get_attention_functions` | Clean. That function is re-invoked on every call by the public wrappers, so patch order doesn't matter. | +| `is_reduce_avg_supported` (FSDP2) | Yes | Not a patch at all — extended `verl_hardware_plugin/engines/fsdp_xpu.py`'s existing `initialize()` to loop over every FSDP2-wrapped submodule, not just the root | See below — this is *better* than the monkeypatch this branch originally planned. | +| `is_reduce_avg_supported` (3 raw `dist.all_reduce(op=AVG)` sites) | Yes, flagged | Patch `torch.distributed.all_reduce` process-wide when `op=AVG` | Robust to import order (attribute lookup at call time), but changes global `ReduceOp.AVG` semantics for the whole process. Real design smell, documented in `patches/reduce_avg_allreduce_patch_xpu.py`. | +| `dist_profiler_cls` (VTune tool selection) | Yes | Patch `DistProfiler.__init__` (mutate the class method, not the name) | Clean — see `patches/dist_profiler_patch_xpu.py` docstring for why patching the method and not reassigning the class name matters. | +| `profiler_markers` (ambient `marked_timer`/`mark_start_range` calls throughout the trainer) | **No** | — | Genuine blocker. `verl/utils/profiler/__init__.py` picks one of three whole modules with an `if/elif/else` **once**, at whatever moment it's first imported by *anything* in the process — there's no per-call dispatch object to patch. Reassigning `verl.utils.profiler.marked_timer` after the fact only helps callers that read that name *after* the patch runs; anything that already did `from verl.utils.profiler import marked_timer` (or `from verl.utils.debug import *`, which re-exports it) earlier keeps its own bound reference. Whether this patch would work in practice depends on winning a race against verl-core's own import order — not a foundation to build on. | +| `dist_profiler_cls` — rollout-server profiler allowlist (vLLM/SGLang/TRT-LLM async servers) | **No** | — | Genuine blocker, same conclusion as `verl-hardware-plugin#26`: `profiler_config = None` and the call that consumes it are in the same method body, with no hook or registry in between. Nothing to intercept short of forking ~40 lines of server constructor per engine. | + +Net: 3 of 4 hooks from #7917 are unnecessary — this plugin already does what +they'd do, with zero core changes. The 4th (`profiler_markers`) turns out to +have the *same* import-order problem as the rollout-server case, which +earlier framing (including an earlier version of this repo's own proposal +doc) had marked as "clean." It isn't. Two things, not one, look like +legitimate small-core-PR asks now — not the 4 hooks #7917 currently has. + +## Why the FSDP2 fix isn't `patches/reduce_avg_fsdp2_patch_xpu.py` + +The first draft of this plan (matching the proposal doc) was to wrap +`verl.utils.fsdp_utils.apply_fsdp2`: call the original, then walk +`model.modules()` and call `set_force_sum_reduction_for_comms(True)`. + +That doesn't actually work reliably: `verl/workers/engine/fsdp/ +transformer_impl.py` does `from verl.utils.fsdp_utils import apply_fsdp2` at +its own import time. Reassigning `verl.utils.fsdp_utils.apply_fsdp2` after +that has already happened has no effect on the name bound inside +`transformer_impl.py` — the exact same class of import-order fragility as +`profiler_markers` above. + +This plugin doesn't need that patch at all, though: `fsdp_xpu.py`'s engine +subclasses already call `set_force_sum_reduction_for_comms` on the +constructed model *after* `super().initialize()` returns — from inside their +own `EngineRegistry`-registered class, which verl-core calls directly (no +patching involved). This branch just extended that existing, already-real +mechanism to loop over every FSDP2-wrapped submodule instead of only the +root (`_force_sum_reduction_on_all_fsdp_modules` in `fsdp_xpu.py`), which is +what `is_reduce_avg_supported`'s FSDP2 wiring in #7917 would have bought — +achieved here with no patch and no hook. + +## Comparison + +| | Files touched outside this plugin | Depends on #7917 merging | +|---|---|---| +| PR #22 (`feature/xpu-vtune-avg-attention`) | 0 (but 4 `PlatformXPU` hook methods are dead code without #7917) | Yes, for all 4 hooks | +| This branch | 0 | No, for 3 of 4. `profiler_markers`'s ambient markers and the rollout-server allowlist remain open — either fork ~40 lines of server code into the plugin, or ask for the two small core changes verl-hardware-plugin#26 already discusses. | + +## Failure policy for the patches + +Every module in `patches/` shares one XPU guard (`patches/_xpu_guard.py`) and +is a hard no-op on a CPU/CUDA/NPU process — the entry point is discovered on +every host, not just Intel ones, and all three patches mutate process-global +state (`torch.distributed.all_reduce`, `verl.utils.attention_utils. +_get_attention_functions`, `DistProfiler.__init__`). Patching any of those on +another vendor's host changes *that* vendor's behavior. + +`apply_all()` splits failures by consequence rather than swallowing them: + +| Patch | On failure, on an XPU host | +|---|---| +| `reduce_avg_allreduce_patch_xpu` | **raise** — oneCCL's `ReduceOp.AVG` is not reliably available; continuing risks a mid-collective abort or wrong gradients | +| `attention_patch_xpu` | **raise** — rmpad attention needs the XPU function set | +| `dist_profiler_patch_xpu` | `logger.exception`, continue — losing `tool: vtune` costs observability, not correctness | + +On a non-XPU host nothing was going to install anyway, so a failure is +debug-logged and never raised. + +## Hardware validation + +Run on 2×B60 (Intel Arc Pro B60 Graphics) via devctl, 2026-09-24, inside the +image built by PR #2's `docker/intel_gpu/Dockerfile.intel_gpu` with this +branch as the build context (verl-core `cbf4f31b`, torch `2.13.0+xpu`, vLLM +`0.29.0`): + +- 85 unit tests pass in-image. +- All three patches install on an XPU host; none install in the same image on + a CPU-only container (`dist.all_reduce` and `DistProfiler.__init__` verified + untouched there). +- `attention_utils._get_attention_functions()` returns the + `verl.utils.npu_flash_attn_utils` set; `profiler.tool: vtune` resolves to + `VtuneProfiler`. +- 2-rank xccl: patched `all_reduce(AVG)` returns the true mean (1.5 for ranks + contributing 1 and 2), `SUM` is unaffected. +- PR #8's premise holds on this image: verl-core `cbf4f31b` *does* define + `attention_utils_module`/`profiler_markers`/`dist_profiler_cls` on + `PlatformBase`, and this plugin overrides none of them — so the hooks are + inert and the patches are what produced the results above. + +One caveat worth recording, since it is easy to misread: a plugin-free +baseline on the same pod showed unpatched xccl computing `ReduceOp.AVG` +*correctly* (both sync and async) for a small 2-rank tensor. That does not +contradict the patch — oneCCL implements AVG on its SYCL-kernel path and not +its scheduler path, and path selection is internal, so a small collective can +land on the working path. It does mean a 2-rank smoke test cannot demonstrate +the failure the patch prevents. + +## Not yet done + +- No end-to-end GRPO training run on this branch; the validation above + exercises the patches directly, not a full trainer loop. +- No unit tests yet for the two capabilities this branch deliberately does + *not* patch (they are argued from source, not exercised). diff --git a/docs/development.md b/docs/development.md index fd2c6207..f20b3925 100644 --- a/docs/development.md +++ b/docs/development.md @@ -577,13 +577,19 @@ class FSDPMyVendorEngineWithLMHead(FSDPEngineWithLMHead): - Force sum reduction if the comm backend doesn't support AVG - Enable/disable specific FSDP features - Apply vendor-specific optimizations + + The wrapped model is `self.module`, assigned by + `FSDPEngine._build_model_optimizer()` (called from `initialize()`) in + verl/workers/engine/fsdp/transformer_impl.py. There is no `self.model` + on these engines -- only `self.model_config` -- so a `hasattr(self.model, + ...)` guard here silently never fires. """ super().initialize() # Example: Force sum-based reduction for custom communication backend # (some backends like xccl don't support ReduceOp.AVG) - if hasattr(self.model, "set_force_sum_reduction_for_comms"): - self.model.set_force_sum_reduction_for_comms(True) + if hasattr(self.module, "set_force_sum_reduction_for_comms"): + self.module.set_force_sum_reduction_for_comms(True) logger.info("Enabled force_sum_reduction for MyVendor comm backend") @@ -607,8 +613,8 @@ class FSDPMyVendorEngineWithValueHead(FSDPEngineWithValueHead): def initialize(self): super().initialize() - if hasattr(self.model, "set_force_sum_reduction_for_comms"): - self.model.set_force_sum_reduction_for_comms(True) + if hasattr(self.module, "set_force_sum_reduction_for_comms"): + self.module.set_force_sum_reduction_for_comms(True) ``` Then register in `verl_hardware_plugin/engines/__init__.py`: @@ -730,8 +736,8 @@ And the corresponding engine: class FSDPXPUEngineWithLMHead(FSDPEngineWithLMHead): def initialize(self): super().initialize() - if hasattr(self.model, "set_force_sum_reduction_for_comms"): - self.model.set_force_sum_reduction_for_comms(True) + if hasattr(self.module, "set_force_sum_reduction_for_comms"): + self.module.set_force_sum_reduction_for_comms(True) ``` --- @@ -849,7 +855,7 @@ Environment variable overrides: **Solution**: In your engine's `initialize()`, enable force sum reduction: ```python -self.model.set_force_sum_reduction_for_comms(True) +self.module.set_force_sum_reduction_for_comms(True) ``` ### Problem: Import errors when plugin is installed on machines without my hardware diff --git a/tests/test_fsdp_xpu_engine.py b/tests/test_fsdp_xpu_engine.py new file mode 100644 index 00000000..c2eaefd7 --- /dev/null +++ b/tests/test_fsdp_xpu_engine.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Regression tests for the Intel XPU FSDP engine's xccl reduce workaround. + +Two things are pinned here: + +1. The attribute name. verl's FSDPEngine assigns the wrapped model to + ``self.module`` (transformer_impl.py::_build_model_optimizer, called from + ``initialize()``); there is no ``self.model`` on these engines, only + ``self.model_config``. + The engine previously guarded on ``hasattr(self.model, ...)``, so the + workaround silently never ran and xccl kept getting ReduceOp.AVG. +2. The recursion. fully_shard() wraps each transformer layer as well as the + root, and each wrapper has its own reduction op, so fixing only the root + leaves per-layer gradient sync on AVG. +""" + +import pytest + +torch = pytest.importorskip("torch") +nn = torch.nn + + +@pytest.fixture(scope="module") +def force_sum_helper(): + fsdp_xpu = pytest.importorskip("verl_hardware_plugin.engines.fsdp_xpu") + return fsdp_xpu._force_sum_reduction_on_all_fsdp_modules + + +class _FakeFSDPModule(nn.Module): + """Stands in for a fully_shard()-wrapped module: has the setter, records calls.""" + + def __init__(self): + super().__init__() + self.forced = None + + def set_force_sum_reduction_for_comms(self, value): + self.forced = value + + +class _FakeFSDPRoot(_FakeFSDPModule): + def __init__(self, n_layers=3): + super().__init__() + self.layers = nn.ModuleList([_FakeFSDPModule() for _ in range(n_layers)]) + # A plain submodule with no setter must simply be skipped, not crash. + self.embed = nn.Linear(4, 4) + + +def test_force_sum_reduction_covers_root_and_every_wrapped_layer(force_sum_helper): + root = _FakeFSDPRoot(n_layers=3) + + force_sum_helper(root) + + assert root.forced is True + assert [layer.forced for layer in root.layers] == [True, True, True] + + +def test_force_sum_reduction_tolerates_a_model_without_the_setter(force_sum_helper): + """FSDP1 / unwrapped models have no setter at all; this must be a no-op.""" + force_sum_helper(nn.Linear(4, 4)) + + +def test_verl_fsdp_engine_exposes_module_not_model(): + """Guards the attribute contract the workaround depends on. + + If a future verl release renames ``module``, this fails here instead of + silently disabling the workaround at runtime. + """ + import inspect + + transformer_impl = pytest.importorskip("verl.workers.engine.fsdp.transformer_impl") + # Scan the whole class, not initialize(): the assignment lives in + # _build_model_optimizer(), which initialize() calls. + class_src = inspect.getsource(transformer_impl.FSDPEngine) + assert "self.module = " in class_src, "verl's FSDPEngine no longer assigns self.module" + assert "self.model = " not in class_src, "verl's FSDPEngine now also assigns self.model; revisit the engine" diff --git a/tests/test_itt_profile_xpu.py b/tests/test_itt_profile_xpu.py new file mode 100644 index 00000000..40ba2d98 --- /dev/null +++ b/tests/test_itt_profile_xpu.py @@ -0,0 +1,111 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Regression tests for the Intel VTune (ITT) profiling backend. + +ITT ranges are a push/pop stack per thread. A range that is pushed but never +popped does not just lose one measurement -- every later range on that thread +nests inside the leaked one, so the trace is unreadable exactly in the runs +where something went wrong and you most want to read it. These tests pin the +push/pop balance when the timed body raises. +""" + +from unittest import mock + +import pytest + +from verl_hardware_plugin.profilers import itt_profile_xpu + + +class _ITTSpy: + """Stand-in for torch.profiler.itt that tracks push/pop depth.""" + + def __init__(self): + self.depth = 0 + self.messages = [] + + def range_push(self, message): + self.messages.append(message) + self.depth += 1 + return self.depth + + def range_pop(self): + self.depth -= 1 + return self.depth + + +@pytest.fixture +def itt(): + spy = _ITTSpy() + with mock.patch.object(itt_profile_xpu, "_itt", spy): + yield spy + + +def test_marked_timer_pops_range_on_success(itt): + timing = {} + with itt_profile_xpu.marked_timer("step", timing): + pass + assert itt.depth == 0 + assert itt.messages == ["step"] + assert "step" in timing + + +def test_marked_timer_pops_range_when_body_raises(itt): + timing = {} + with pytest.raises(ValueError, match="boom"): + with itt_profile_xpu.marked_timer("step", timing): + raise ValueError("boom") + assert itt.depth == 0, "a raising body must not leave the ITT range open" + + +def test_marked_timer_propagates_the_original_exception(itt): + """The timing exception must survive the finally, not be replaced by it.""" + with pytest.raises(KeyError): + with itt_profile_xpu.marked_timer("step", {}): + raise KeyError("original") + + +def test_annotate_pops_range_when_decorated_function_raises(itt): + profiler = itt_profile_xpu.VtuneProfiler(rank=0, config=None, tool_config=None) + + @profiler.annotate(message="rollout") + def boom(): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + boom() + assert itt.depth == 0, "a raising decorated function must not leave the ITT range open" + assert itt.messages == ["rollout"] + + +def test_annotate_returns_value_and_pops_range_on_success(itt): + profiler = itt_profile_xpu.VtuneProfiler(rank=0, config=None, tool_config=None) + + @profiler.annotate() + def add(a, b): + return a + b + + assert add(1, 2) == 3 + assert itt.depth == 0 + assert itt.messages == ["add"] + + +def test_annotate_stops_discrete_collection_when_function_raises(itt): + """With discrete=True the platform profiler is started per call; it must be + stopped on the error path too, or the collector keeps running forever.""" + profiler = itt_profile_xpu.VtuneProfiler(rank=0, config=None, tool_config=None) + profiler.discrete = True + + platform = mock.MagicMock() + with mock.patch.object(itt_profile_xpu, "get_platform", return_value=platform): + + @profiler.annotate() + def boom(): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError): + boom() + + platform.profiler_start.assert_called_once() + platform.profiler_stop.assert_called_once() + assert itt.depth == 0 diff --git a/tests/test_xpu_patches.py b/tests/test_xpu_patches.py new file mode 100644 index 00000000..049d8c3f --- /dev/null +++ b/tests/test_xpu_patches.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Regression tests for verl_hardware_plugin.patches (Intel XPU monkeypatches). + +Each patch in that package mutates global verl-core/torch state, so the two +things worth pinning down are: + +1. On a non-XPU host every patch is a no-op, and nothing global is touched. + (A CPU/CUDA/NPU process still imports this plugin -- the entry point is + discovered on every host.) +2. On an XPU host a patch that fails to install is not swallowed: a + correctness-critical one raises, a feature one is logged loudly. + +These run on a CPU-only host; XPU presence is simulated by patching the shared +guard, so no Intel GPU is required. +""" + +import logging +from contextlib import ExitStack +from unittest import mock + +import pytest + +from verl_hardware_plugin import patches +from verl_hardware_plugin.patches import ( + attention_patch_xpu, + dist_profiler_patch_xpu, + reduce_avg_allreduce_patch_xpu, +) + +PATCH_MODULES = (attention_patch_xpu, dist_profiler_patch_xpu, reduce_avg_allreduce_patch_xpu) + + +@pytest.fixture(autouse=True) +def _reset_applied_flags(): + """Patches are idempotent via a module-level _applied flag; reset it per test.""" + originals = {m: m._applied for m in PATCH_MODULES} + for m in PATCH_MODULES: + m._applied = False + yield + for m, value in originals.items(): + m._applied = value + + +def _force_xpu_availability(available: bool): + """Patch the shared guard as each module sees it (all import it by name).""" + stack = ExitStack() + for target in (patches, *PATCH_MODULES): + stack.enter_context(mock.patch.object(target, "xpu_available", return_value=available)) + return stack + + +@pytest.fixture +def no_xpu(): + with _force_xpu_availability(False): + yield + + +@pytest.fixture +def fake_xpu(): + with _force_xpu_availability(True): + yield + + +@pytest.mark.parametrize("module", PATCH_MODULES, ids=lambda m: m.__name__.rsplit(".", 1)[-1]) +def test_patch_is_noop_without_xpu(module, no_xpu): + """Every patch module must check for XPU before mutating anything global. + + dist_profiler_patch_xpu originally did not, so importing this plugin on a + CUDA host rerouted that host's `profiler.tool: vtune` to the Intel ITT + profiler. + """ + module.apply() + assert module._applied is False, f"{module.__name__} installed itself on a non-XPU host" + + +def test_dist_profiler_untouched_without_xpu(no_xpu): + """The concrete consequence of the guard above, checked on the real class.""" + from verl.utils.profiler.profile import DistProfiler + + before = DistProfiler.__init__ + dist_profiler_patch_xpu.apply() + assert DistProfiler.__init__ is before + + +def test_all_reduce_untouched_without_xpu(no_xpu): + import torch.distributed as dist + + before = dist.all_reduce + reduce_avg_allreduce_patch_xpu.apply() + assert dist.all_reduce is before + + +def test_apply_all_is_quiet_on_non_xpu_host(no_xpu, caplog): + """A failing patch on a non-XPU host is uninteresting: debug-log, no raise.""" + with mock.patch.object(attention_patch_xpu, "apply", side_effect=ImportError("boom")): + with caplog.at_level(logging.DEBUG, logger="verl_hardware_plugin.patches"): + patches.apply_all() + assert not [r for r in caplog.records if r.levelno >= logging.WARNING] + + +def test_apply_all_raises_when_required_patch_fails_on_xpu(fake_xpu): + """A correctness-critical patch that did not install must not be swallowed. + + Without this, an XPU run continues with real ReduceOp.AVG on xccl and + silently produces wrong gradients. + """ + with mock.patch.object(reduce_avg_allreduce_patch_xpu, "apply", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError, match="reduce_avg_allreduce_patch_xpu"): + patches.apply_all() + + +def test_apply_all_reports_but_survives_optional_patch_failure_on_xpu(fake_xpu, caplog): + """The vtune patch is observability, not correctness: log loudly, keep going.""" + with mock.patch.object(dist_profiler_patch_xpu, "apply", side_effect=RuntimeError("boom")): + with mock.patch.object(attention_patch_xpu, "apply"): + with mock.patch.object(reduce_avg_allreduce_patch_xpu, "apply"): + with caplog.at_level(logging.DEBUG, logger="verl_hardware_plugin.patches"): + patches.apply_all() + + errors = [r for r in caplog.records if r.levelno >= logging.ERROR] + assert errors, "an optional patch failing on an XPU host must be visible above debug level" + assert "dist_profiler_patch_xpu" in errors[0].getMessage() diff --git a/verl_hardware_plugin/__init__.py b/verl_hardware_plugin/__init__.py index 90ee349b..90b7e4ee 100644 --- a/verl_hardware_plugin/__init__.py +++ b/verl_hardware_plugin/__init__.py @@ -14,6 +14,7 @@ import os from verl_hardware_plugin.engines import register_all_engines +from verl_hardware_plugin.patches import apply_all as apply_all_patches from verl_hardware_plugin.platforms import register_all_platforms from verl_hardware_plugin.profilers import register_all_profiles @@ -23,5 +24,6 @@ register_all_platforms() register_all_engines() register_all_profiles() +apply_all_patches() logger.info("verl-hardware-plugin loaded successfully") diff --git a/verl_hardware_plugin/engines/fsdp_xpu.py b/verl_hardware_plugin/engines/fsdp_xpu.py index 6c7397b1..95c0f1e7 100644 --- a/verl_hardware_plugin/engines/fsdp_xpu.py +++ b/verl_hardware_plugin/engines/fsdp_xpu.py @@ -13,6 +13,21 @@ followed by manual division, which is functionally equivalent but compatible with xccl. +Why not patch verl-core's apply_fsdp2() instead? + verl/workers/engine/fsdp/transformer_impl.py does + `from verl.utils.fsdp_utils import apply_fsdp2` at its own import time, so a + plugin that reassigns `verl.utils.fsdp_utils.apply_fsdp2` after that module + has already imported the name would have no effect on this call site -- + the same import-order fragility as verl core's profiler-marker dispatch + (see docs/design/ for that one). This engine sidesteps the problem + entirely: it does not patch anything. It calls + set_force_sum_reduction_for_comms() on the already-constructed model + *after* super().initialize() (which internally calls the real + apply_fsdp2()) returns, from inside the plugin's own EngineRegistry + subclass. That is ordinary method-override behavior, not a monkeypatch, + so it is robust regardless of when/how anything else imported + apply_fsdp2. + Registration: @EngineRegistry.register(device="xpu", vendor="intel") This means verl will automatically select this engine when: @@ -38,6 +53,28 @@ logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) +def _force_sum_reduction_on_all_fsdp_modules(root) -> None: + """Apply set_force_sum_reduction_for_comms(True) to root and every nested + FSDP2-wrapped submodule, not just the root. + + fully_shard() wraps the transformer layers individually as well as the + root model (verl/utils/fsdp_utils.py::apply_fsdp2), so each of those + per-layer wrappers has its own communication group that also defaults to + ReduceOp.AVG. Only fixing the root leaves per-layer gradient sync using + AVG, which xccl silently cannot execute correctly. + """ + if hasattr(root, "set_force_sum_reduction_for_comms"): + root.set_force_sum_reduction_for_comms(True) + count = 1 + for submodule in root.modules(): + if submodule is root: + continue + if hasattr(submodule, "set_force_sum_reduction_for_comms"): + submodule.set_force_sum_reduction_for_comms(True) + count += 1 + logger.info("Enabled force_sum_reduction_for_comms on %d FSDP module(s) for XPU", count) + + @EngineRegistry.register(model_type="language_model", backend=["fsdp", "fsdp2"], device="xpu", vendor="intel") class FSDPXPUEngineWithLMHead(FSDPEngineWithLMHead): """FSDP Engine for Intel XPU with xccl communication backend. @@ -59,15 +96,14 @@ def __init__( def initialize(self): """Initialize the FSDP model, then apply XPU-specific workarounds. - The key workaround: force sum-based gradient reduction. - This is needed because xccl does not support ReduceOp.AVG. - The FSDP wrapper will use SUM + manual division instead. + The key workaround: force sum-based gradient reduction, root module + and every per-layer FSDP2 wrapper. This is needed because xccl does + not support ReduceOp.AVG. The FSDP wrapper will use SUM + manual + division instead. """ super().initialize() # xccl does not support ReduceOp.AVG; force sum-based reduction - if hasattr(self.model, "set_force_sum_reduction_for_comms"): - self.model.set_force_sum_reduction_for_comms(True) - logger.info("Enabled force_sum_reduction_for_comms for XPU") + _force_sum_reduction_on_all_fsdp_modules(self.module) @EngineRegistry.register(model_type="value_model", backend=["fsdp", "fsdp2"], device="xpu", vendor="intel") @@ -91,5 +127,4 @@ def __init__( def initialize(self): """Initialize the FSDP value model, then apply xccl workaround.""" super().initialize() - if hasattr(self.model, "set_force_sum_reduction_for_comms"): - self.model.set_force_sum_reduction_for_comms(True) + _force_sum_reduction_on_all_fsdp_modules(self.module) diff --git a/verl_hardware_plugin/patches/__init__.py b/verl_hardware_plugin/patches/__init__.py new file mode 100644 index 00000000..03d68488 --- /dev/null +++ b/verl_hardware_plugin/patches/__init__.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Plugin-side monkeypatches for Intel XPU, applied with zero changes to +verl-core and with no dependency on any new PlatformBase hook. + +This package exists as the "how far can pure monkeypatch actually go" +counterpart to the PlatformBase-hook approach used elsewhere in this plugin +(is_reduce_avg_supported/attention_utils_module/profiler_markers/ +dist_profiler_cls, gated behind the still-unmerged verl-project/verl#7917). +See docs/design/xpu-monkeypatch-experiment.md for the per-capability verdict +and why two capabilities are NOT here (they are genuine blockers, not solved +by patching harder). + +Every patch module is independently guarded by the shared XPU-availability +check in ``_xpu_guard`` and sets its own ``_applied`` flag, so apply_all() is +idempotent and importing this package is safe even when another vendor's +plugin (or no accelerator at all) is active in the same process. + +Failure handling + On a non-XPU process there is nothing to install, so a failure is not + interesting and is logged at debug. On an XPU process a patch that fails + to install means an advertised capability is silently absent, which is + exactly the kind of thing that later surfaces as a wrong number or an + unexplained collective hang -- so those are never swallowed: + + - REQUIRED patches guard numerical correctness (oneCCL's ReduceOp.AVG is + not reliably available -- see reduce_avg_allreduce_patch_xpu's docstring + for why "not reliably" rather than "not at all"; rmpad attention needs + the XPU function set). If one of these cannot be installed, training + would produce wrong results or abort far from the cause, so apply_all() + re-raises as RuntimeError. + - Non-required patches are feature wiring (`profiler.tool: vtune` + selection). Losing one degrades observability but not correctness, so it + is reported with logger.exception and training continues. +""" + +import importlib +import logging +import os + +from ._xpu_guard import xpu_available + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +# (module attribute name, required-for-correctness) +_PATCH_MODULES = ( + ("reduce_avg_allreduce_patch_xpu", True), + ("attention_patch_xpu", True), + ("dist_profiler_patch_xpu", False), +) + + +def apply_all() -> None: + """Apply every XPU monkeypatch. Safe to call multiple times or on non-XPU hosts. + + Raises: + RuntimeError: on an XPU host, if a correctness-critical patch could not + be installed. Never raises on a non-XPU host. + """ + on_xpu = xpu_available() + + for name, required in _PATCH_MODULES: + qualname = f"{__name__}.{name}" + try: + # Imported here, not at module scope, so that a module-level import + # failure is subject to the same policy as an apply() failure. + importlib.import_module(f".{name}", __name__).apply() + except Exception as e: + if not on_xpu: + # Nothing was going to be installed on this host anyway. + logger.debug("%s.apply() skipped on non-XPU host: %s", qualname, e) + elif required: + raise RuntimeError( + f"{qualname}.apply() failed on an Intel XPU host. This patch is required for correct " + f"results (oneCCL's ReduceOp.AVG is not reliably available / rmpad attention needs the " + f"XPU function set); continuing would silently produce wrong numbers." + ) from e + else: + logger.exception( + "%s.apply() failed on an Intel XPU host; the capability it provides is unavailable", + qualname, + ) diff --git a/verl_hardware_plugin/patches/_xpu_guard.py b/verl_hardware_plugin/patches/_xpu_guard.py new file mode 100644 index 00000000..af1b0e2c --- /dev/null +++ b/verl_hardware_plugin/patches/_xpu_guard.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Shared Intel XPU availability check for the patch modules. + +Every module in this package mutates global verl-core / torch state, so each +one must be a no-op on a CPU/CUDA/NPU process that happens to import this +plugin (the entry point is discovered on every host, not just Intel ones). +One helper, used by all of them, so the guard cannot drift between modules. +""" + + +def xpu_available() -> bool: + """True only when this process actually has a usable Intel XPU device.""" + import torch + + return hasattr(torch, "xpu") and torch.xpu.is_available() diff --git a/verl_hardware_plugin/patches/attention_patch_xpu.py b/verl_hardware_plugin/patches/attention_patch_xpu.py new file mode 100644 index 00000000..3191cbc9 --- /dev/null +++ b/verl_hardware_plugin/patches/attention_patch_xpu.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Monkeypatch to route Intel XPU through the NPU flash-attention shim for +rmpad attention padding, with zero verl-core changes. + +Dispatch shape that makes this safe (verl/utils/attention_utils.py): + verl.utils.attention_utils.index_first_axis/pad_input/rearrange/unpad_input + are stable public wrappers. Each one calls the private + `_get_attention_functions()` fresh on *every* invocation -- it is not + cached at import time. That means patching `_get_attention_functions` + itself takes effect for every caller regardless of whether that caller + did `from verl.utils.attention_utils import pad_input` before or after + this patch runs. This is the one hook where import order genuinely does + not matter. +""" + +import logging +import os + +from ._xpu_guard import xpu_available + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +_applied = False + + +def apply() -> None: + global _applied + if _applied: + return + if not xpu_available(): + return + + import verl.utils.attention_utils as attention_utils + + original_get_attention_functions = attention_utils._get_attention_functions + + def _patched_get_attention_functions(): + if xpu_available(): + from verl.utils.npu_flash_attn_utils import index_first_axis, pad_input, rearrange, unpad_input + + attention_utils._index_first_axis = index_first_axis + attention_utils._pad_input = pad_input + attention_utils._rearrange = rearrange + attention_utils._unpad_input = unpad_input + return index_first_axis, pad_input, rearrange, unpad_input + return original_get_attention_functions() + + attention_utils._get_attention_functions = _patched_get_attention_functions + _applied = True + logger.info("[verl_hardware_plugin] Patched attention_utils._get_attention_functions for XPU") diff --git a/verl_hardware_plugin/patches/dist_profiler_patch_xpu.py b/verl_hardware_plugin/patches/dist_profiler_patch_xpu.py new file mode 100644 index 00000000..6ade9bcb --- /dev/null +++ b/verl_hardware_plugin/patches/dist_profiler_patch_xpu.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Monkeypatch DistProfiler.__init__ so `profiler.tool: vtune` resolves to +VtuneProfiler, with zero verl-core changes and no dependency on +PlatformBase.dist_profiler_cls(). + +Why patching the class method (not reassigning the name `DistProfiler`) +matters: + verl/utils/profiler/profile.py::DistProfiler.__init__ has a hardcoded + if/elif chain on `self._tool` ("nsys", "npu", "torch", "torch_memory", + "precision_debugger") with no registry -- there is no hook to add a + "vtune" branch to. But `DistProfiler` is a class object shared by every + caller regardless of how or when they imported it (`from + verl.utils.profiler import DistProfiler` vs `verl.utils.profiler.profile. + DistProfiler`) -- reassigning its `__init__` attribute mutates that one + shared object, so it takes effect for every existing and future caller. + That is different from e.g. reduce_avg's `apply_fsdp2`, which is a plain + function imported by name into another module's namespace before this + plugin gets a chance to patch it. + +Wrap-then-fixup is safe here because DistProfiler.__init__ does not have side +effects beyond constructing `self._impl`: run the real __init__, then swap +`self._impl` for VtuneProfiler if the real __init__ fell back to +`_NoOpProfiler` for a tool name it doesn't recognize and that name is +"vtune". Nothing destructive happens in the original that a post-hoc fixup +could not undo. + +This does NOT solve verl-project/verl-hardware-plugin#26's +`profiler_markers` hook: the ambient `marked_timer()`/`mark_start_range()` +calls scattered through verl-core's own trainer code (verl/utils/profiler/ +__init__.py's import-time module selection) are a separate, genuine blocker +-- see docs/design/xpu-monkeypatch-experiment.md. What this patch gets you: +`profiler.tool: vtune` correctly selects VtuneProfiler, and its own +start()/stop()/annotate() emit real ITT ranges (VtuneProfiler calls its own +local mark_start_range/mark_end_range, not the ambient ones). What it does +NOT get you: the rest of the trainer's own `marked_timer()` calls elsewhere +in verl-core emitting ITT ranges too. +""" + +import logging +import os + +from ._xpu_guard import xpu_available + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +_applied = False + + +def apply() -> None: + global _applied + if _applied: + return + # DistProfiler is a process-wide shared class object, so patching it on a + # CPU/CUDA/NPU process would reroute *that* platform's `tool: vtune` to the + # Intel ITT implementation. Same guard as every other patch in this package. + if not xpu_available(): + return + + from verl.utils.profiler.profile import DistProfiler, _NoOpProfiler + + original_init = DistProfiler.__init__ + + def _patched_init(self, rank, config=None, tool_config=None, save_file_prefix=None, **kwargs): + original_init(self, rank, config=config, tool_config=tool_config, save_file_prefix=save_file_prefix, **kwargs) + + if self._tool == "vtune" and isinstance(self._impl, _NoOpProfiler): + from verl_hardware_plugin.profilers.itt_profile_xpu import VtuneProfiler + + self._impl = VtuneProfiler(rank=rank, config=self.config, tool_config=self.tool_config, **kwargs) + + DistProfiler.__init__ = _patched_init + _applied = True + logger.info("[verl_hardware_plugin] Patched DistProfiler.__init__: +vtune") diff --git a/verl_hardware_plugin/patches/reduce_avg_allreduce_patch_xpu.py b/verl_hardware_plugin/patches/reduce_avg_allreduce_patch_xpu.py new file mode 100644 index 00000000..e83073f4 --- /dev/null +++ b/verl_hardware_plugin/patches/reduce_avg_allreduce_patch_xpu.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Monkeypatch torch.distributed.all_reduce(op=AVG) -> SUM + manual divide, +process-wide, for Intel XPU (xccl does not implement ReduceOp.AVG). + +This is the one patch in this package worth calling "doable, but with a real +design smell" rather than "clean": + +- verl/workers/engine_workers.py, verl/trainer/sft_trainer.py, and + verl/utils/profiler/performance.py::reduce_timing each call + `torch.distributed.all_reduce(x, op=ReduceOp.AVG, ...)` as a literal inline + expression -- there is no per-site wrapper or dispatch object to patch, the + way there is for attention padding. +- Patching `torch.distributed.all_reduce` itself is robust to import order + (every call site does `dist.all_reduce(...)`, an attribute lookup on the + `torch.distributed` module performed at call time, not at each file's own + import time), so this does take effect regardless of load order. +- The cost is real: this changes the process-wide meaning of + `ReduceOp.AVG` for every call site in the process, XPU-related or not. + Anyone reading verl-core's source sees real `ReduceOp.AVG` semantics and + silently gets SUM+divide instead. That is a legitimate argument for fixing + this in verl-core instead (verl-project/verl#7917's + `is_reduce_avg_supported()` hook covers exactly these 3 call sites) -- this + patch exists to show it is *possible* without core changes, not to claim + it is the better design. + +async_op=True is intentionally unsupported: correctly dividing the result +requires the collective to have already completed, which async_op explicitly +defers. Call sites using async_op=True with op=AVG fall through to the +original (broken-on-xccl) behavior rather than silently producing a wrong +answer with no error. +""" + +import logging +import os + +from ._xpu_guard import xpu_available + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +_applied = False + + +def apply() -> None: + global _applied + if _applied: + return + if not xpu_available(): + return + + import torch.distributed as dist + + original_all_reduce = dist.all_reduce + + def _patched_all_reduce(tensor, op=dist.ReduceOp.SUM, group=None, async_op=False): + if op != dist.ReduceOp.AVG or async_op: + return original_all_reduce(tensor, op=op, group=group, async_op=async_op) + + world_size = dist.get_world_size(group=group) + result = original_all_reduce(tensor, op=dist.ReduceOp.SUM, group=group, async_op=False) + tensor.div_(world_size) + return result + + dist.all_reduce = _patched_all_reduce + _applied = True + logger.info("[verl_hardware_plugin] Patched torch.distributed.all_reduce(op=AVG) for XPU/xccl") diff --git a/verl_hardware_plugin/profilers/itt_profile_xpu.py b/verl_hardware_plugin/profilers/itt_profile_xpu.py new file mode 100644 index 00000000..ec694f7b --- /dev/null +++ b/verl_hardware_plugin/profilers/itt_profile_xpu.py @@ -0,0 +1,157 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Intel VTune (ITT) profiling backend for verl. + +Wired in via verl_hardware_plugin.patches.dist_profiler_patch_xpu, which +monkeypatches DistProfiler.__init__ to select VtuneProfiler when +`profiler.tool: vtune` -- not via a PlatformBase hook. See that module's +docstring for what this patch does and does not cover (ambient +marked_timer()/mark_start_range() calls elsewhere in verl-core's trainer +code are NOT covered; only VtuneProfiler's own start/stop/annotate are). +""" + +import functools +from contextlib import contextmanager +from typing import Callable, Optional + +import torch.profiler.itt as _itt + +from verl.plugin.platform import get_platform +from verl.utils.profiler.config import NsightToolConfig +from verl.utils.profiler.profile import DistProfiler, ProfilerConfig + + +def mark_start_range( + message: Optional[str] = None, + color: Optional[str] = None, + domain: Optional[str] = None, + category: Optional[str] = None, +) -> int: + """Push an ITT range onto the stack. Returns the stack depth (used as range_id). + + color/domain/category are accepted for API compatibility but ignored — ITT does not + support them. + """ + return _itt.range_push(message or "") + + +def mark_end_range(range_id: int) -> None: + """Pop the innermost ITT range. range_id is accepted for API compatibility.""" + _itt.range_pop() + + +def mark_annotate( + message: Optional[str] = None, + color: Optional[str] = None, + domain: Optional[str] = None, + category: Optional[str] = None, +) -> Callable: + """Decorate a function to wrap its execution in an ITT range.""" + + def decorator(func): + profile_message = message or func.__name__ + + @functools.wraps(func) + def wrapper(*args, **kwargs): + _itt.range_push(profile_message) + try: + return func(*args, **kwargs) + finally: + _itt.range_pop() + + return wrapper + + return decorator + + +@contextmanager +def marked_timer( + name: str, + timing_raw: dict[str, float], + color: Optional[str] = None, + domain: Optional[str] = None, + category: Optional[str] = None, +): + """Context manager for timing with ITT ranges (Intel VTune). + + Measures execution time, accumulates into timing_raw, and emits an ITT range + visible in VTune when the process runs under a collector. + + The range is popped in a ``finally`` so a raising body cannot leave it open. + (verl-core's nvtx equivalent does not do this; an exception there leaks an + NVTX range and every later range on the thread nests inside it, which makes + the resulting trace unreadable exactly when you most want to read it.) + """ + mark_range = mark_start_range(message=name) + from verl.utils.profiler.performance import _timer + + try: + yield from _timer(name, timing_raw) + finally: + mark_end_range(mark_range) + + +class VtuneProfiler(DistProfiler): + """Intel VTune profiler. Installed in a worker to control ITT ranges. + + Unlike Nsight (which has a process-level start/stop API), VTune attaches + externally as a collector and observes range_push/range_pop events, so + PlatformXPU.profiler_start/profiler_stop are no-ops by design — the ranges + emitted by mark_start_range/mark_end_range are the real signal. + + Because those two hooks are no-ops, ``tool_config.discrete`` has no effect on + XPU: both values behave identically. It is read only to keep the DistProfiler + contract, so a missing tool_config is not an error here. + """ + + def __init__(self, rank: int, config: Optional[ProfilerConfig], tool_config: Optional[NsightToolConfig], **kwargs): + if not config: + config = ProfilerConfig(ranks=[]) + # verl core has no `tool_config.vtune` schema entry, so `tool_config.get("vtune")` + # resolves to None and DistProfiler then substitutes the whole tool_config mapping + # (verl/utils/profiler/profile.py: `if tool_config is None: tool_config = config.tool_config`). + # That mapping is truthy but has no `discrete`, so read it defensively — the same way + # core itself does. `discrete` is a no-op on XPU anyway, so False is always correct. + self.discrete: bool = getattr(tool_config, "discrete", False) + + def start(self, **kwargs): + if not self.discrete: + get_platform().profiler_start() + + def stop(self): + if not self.discrete: + get_platform().profiler_stop() + + def step(self): + return + + def annotate( + self, + message: Optional[str] = None, + color: Optional[str] = None, + domain: Optional[str] = None, + category: Optional[str] = None, + **kwargs_outer, + ) -> Callable: + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs_inner): + profile_name = message or func.__name__ + + if self.discrete: + get_platform().profiler_start() + mark_range = mark_start_range(message=profile_name) + + try: + return func(*args, **kwargs_inner) + finally: + # Pop the range even if func raised, or every later range on + # this thread nests inside the leaked one. + mark_end_range(mark_range) + if self.discrete: + get_platform().profiler_stop() + + return wrapper + + return decorator