diff --git a/tests/test_plugin_registration.py b/tests/test_plugin_registration.py index b2288b9d..9682a7c5 100644 --- a/tests/test_plugin_registration.py +++ b/tests/test_plugin_registration.py @@ -689,5 +689,148 @@ def test_raises_on_whitelist_and_blacklist(self): may_enable_flag_gems(phase="training") +class TestReduceAvgPatchWiring: + """The reduce_avg monkeypatch must fire on platform *selection*, not on + mere XPU hardware/SDK *presence*. + + verl.plugin.platform.platform_manager._create_platform() only constructs + PlatformXPU() once platform detection (VERL_PLATFORM=intel, or + auto-detection) has actually picked "intel" for this process. Applying + the patch from the shared verl_hardware_plugin/__init__.py instead would + fire on any host where XPU merely happens to be importable, even when a + different platform is selected for the run -- e.g. VERL_PLATFORM=nvidia + explicitly set on a mixed host would still get its process-wide + torch.distributed.all_reduce(op=AVG) semantics changed for no reason. + """ + + def test_platform_xpu_init_applies_patches(self): + from verl_hardware_plugin.platforms.platform_xpu import PlatformXPU + + with mock.patch("verl_hardware_plugin.patches.xpu.reduce_avg_allreduce_patch.apply") as fake_apply: + PlatformXPU() + + fake_apply.assert_called_once() + + def test_plugin_init_does_not_bind_reduce_avg_patch(self): + import verl_hardware_plugin + + assert not hasattr(verl_hardware_plugin, "apply_all_patches") + assert not hasattr(verl_hardware_plugin, "reduce_avg_allreduce_patch_xpu") + + +class TestReduceAvgAllReducePatch: + """torch.distributed.all_reduce(op=AVG) -> SUM + manual divide, scoped to + xccl process groups only -- a gloo/nccl group in the same process (e.g. a + CPU-only Ray actor's coordination group) must keep native AVG. + + _applied is process-global module state (the same property that makes + this patch risky in production -- see the module's own docstring), so + each test must save/restore torch.distributed.all_reduce and reset the + flag, or tests would leak into each other. + """ + + @pytest.fixture(autouse=True) + def _reset_patch_state(self): + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + original_all_reduce = dist.all_reduce + patch_mod._applied = False + yield + dist.all_reduce = original_all_reduce + patch_mod._applied = False + + def test_noop_when_xpu_unavailable(self): + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + before = dist.all_reduce + with mock.patch.object(patch_mod, "_xpu_available", return_value=False): + patch_mod.apply() + + assert dist.all_reduce is before + + def test_avg_becomes_sum_plus_divide(self): + import torch + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + fake_original = mock.MagicMock() + tensor = torch.tensor([4.0]) + with mock.patch.object(patch_mod, "_xpu_available", return_value=True): + with mock.patch.object(dist, "all_reduce", fake_original): + patch_mod.apply() + patched = dist.all_reduce + assert patched is not fake_original + + with mock.patch.object(dist, "get_backend", return_value="xccl"): + with mock.patch.object(dist, "get_world_size", return_value=4): + patched(tensor, op=dist.ReduceOp.AVG) + + fake_original.assert_called_once_with(tensor, op=dist.ReduceOp.SUM, group=None, async_op=False) + assert tensor.item() == 1.0 + + def test_non_xccl_backend_passes_through_even_with_avg(self): + """A gloo/nccl group in the same process must keep native AVG untouched.""" + import torch + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + fake_original = mock.MagicMock() + tensor = torch.tensor([4.0]) + with mock.patch.object(patch_mod, "_xpu_available", return_value=True): + with mock.patch.object(dist, "all_reduce", fake_original): + patch_mod.apply() + with mock.patch.object(dist, "get_backend", return_value="gloo"): + dist.all_reduce(tensor, op=dist.ReduceOp.AVG, group="cpu_group") + + fake_original.assert_called_once_with(tensor, op=dist.ReduceOp.AVG, group="cpu_group", async_op=False) + + def test_non_avg_op_passes_through_unchanged(self): + import torch + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + fake_original = mock.MagicMock() + tensor = torch.tensor([4.0]) + with mock.patch.object(patch_mod, "_xpu_available", return_value=True): + with mock.patch.object(dist, "all_reduce", fake_original): + patch_mod.apply() + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group="g") + + fake_original.assert_called_once_with(tensor, op=dist.ReduceOp.SUM, group="g", async_op=False) + + def test_async_avg_falls_through_unpatched(self): + import torch + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + fake_original = mock.MagicMock() + tensor = torch.tensor([4.0]) + with mock.patch.object(patch_mod, "_xpu_available", return_value=True): + with mock.patch.object(dist, "all_reduce", fake_original): + patch_mod.apply() + dist.all_reduce(tensor, op=dist.ReduceOp.AVG, async_op=True) + + fake_original.assert_called_once_with(tensor, op=dist.ReduceOp.AVG, group=None, async_op=True) + + def test_idempotent(self): + import torch.distributed as dist + + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch as patch_mod + + with mock.patch.object(patch_mod, "_xpu_available", return_value=True): + patch_mod.apply() + patched_once = dist.all_reduce + patch_mod.apply() + assert dist.all_reduce is patched_once + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/verl_hardware_plugin/engines/fsdp_xpu.py b/verl_hardware_plugin/engines/fsdp_xpu.py index 6c7397b1..77f58251 100644 --- a/verl_hardware_plugin/engines/fsdp_xpu.py +++ b/verl_hardware_plugin/engines/fsdp_xpu.py @@ -7,11 +7,34 @@ (e.g., force sum reduction for xccl backend). Why is this engine needed? - The Intel xccl collective communication backend does NOT support - ReduceOp.AVG for allreduce operations. FSDP's gradient synchronization + Intel's oneCCL (xccl) collective library implements ReduceOp.AVG on its + SYCL-kernel execution path, but not on its scheduler path -- and which + path a given collective takes is selected internally by oneCCL, not + capability-aware, so an AVG request can abort mid-collective + (`average operation is not supported for the scheduler path`) rather + than reliably succeed or reliably fail. FSDP's gradient synchronization normally uses AVG for efficiency. This engine forces sum-based reduction - followed by manual division, which is functionally equivalent but - compatible with xccl. + followed by manual division, which is functionally equivalent and, + unlike AVG, not subject to that path-selection abort. (Fix tracked for + oneCCL 2022.2 / torch 2.15; not yet available on the stack this plugin + targets. A separate, narrower double-division bug on very small + messages, intel/torch-xpu-ops#3020, is fixed on that same stack and is + not what this workaround is for.) + +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 that motivates verl core's own + profiler-marker dispatch design. 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") @@ -38,6 +61,29 @@ 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 is subject to oneCCL's unreliable AVG path selection (see + module docstring) instead of the guaranteed-correct SUM+divide path. + """ + 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 +105,16 @@ 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 oneCCL's + AVG path selection isn't capability-aware and can abort mid-collective + (see module docstring). 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") + # oneCCL's AVG path selection isn't capability-aware and can abort; + # force sum-based reduction instead (see module docstring) + _force_sum_reduction_on_all_fsdp_modules(self.module) @EngineRegistry.register(model_type="value_model", backend=["fsdp", "fsdp2"], device="xpu", vendor="intel") @@ -91,5 +138,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..397f1c84 --- /dev/null +++ b/verl_hardware_plugin/patches/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Plugin-side monkeypatches, one module per capability, applied only from +the platform class that needs them -- not from this package's `__init__.py` +and not from this plugin's shared top-level `verl_hardware_plugin/__init__.py`. + +Each `PlatformXXX()` is only constructed once `verl.plugin.platform. +platform_manager._create_platform()` has actually selected that platform for +the process, so a platform's own `__init__` is the right place to apply its +patches. Applying a patch from a shared location instead would fire on +any host where the corresponding hardware/SDK merely happens to be +importable, regardless of which platform verl actually selects for the run. +""" diff --git a/verl_hardware_plugin/patches/xpu/reduce_avg_allreduce_patch.py b/verl_hardware_plugin/patches/xpu/reduce_avg_allreduce_patch.py new file mode 100644 index 00000000..683dc179 --- /dev/null +++ b/verl_hardware_plugin/patches/xpu/reduce_avg_allreduce_patch.py @@ -0,0 +1,89 @@ +# 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, +scoped to xccl process groups, for Intel XPU. + +Why this is needed: Intel's oneCCL (xccl) implements ReduceOp.AVG on its +SYCL-kernel execution path, but not on its scheduler path. Which path a +given collective takes is selected internally by oneCCL and isn't +capability-aware, so a caller can't tell in advance whether an AVG request +will succeed or abort mid-collective (`average operation is not supported +for the scheduler path`). Forcing sum-based reduction + a manual divide is +the only way to guarantee correctness until that dispatch becomes +capability-aware -- fix tracked for oneCCL 2022.2 / torch 2.15, not yet +available on the stack this plugin targets. (A separate, narrower +double-division bug on very small messages, intel/torch-xpu-ops#3020, is +fixed on that same stack and is not what this patch works around.) + +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: this still intercepts a public torch API process-wide -- once + applied, every `dist.all_reduce(op=AVG)` call in the process is routed + through this wrapper, XPU-related or not. It is scoped to xccl groups via + `dist.get_backend(group) != "xccl"`, so a gloo/nccl group in the same + process (e.g. a CPU-only Ray actor's coordination group alongside a GPU + worker's xccl group) is untouched and keeps native `AVG` -- but anyone + reading verl-core's source still can't see, from that source alone, that + an xccl group's `ReduceOp.AVG` silently becomes SUM+divide. 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 behavior -- subject to the same oneCCL path-selection risk this +patch exists to avoid -- rather than silently producing a wrong answer with +no error. +""" + +import logging +import os + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +_applied = False + + +def _xpu_available() -> bool: + import torch + + return hasattr(torch, "xpu") and torch.xpu.is_available() + + +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 or dist.get_backend(group) != "xccl": + 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 xccl process groups") diff --git a/verl_hardware_plugin/platforms/platform_xpu.py b/verl_hardware_plugin/platforms/platform_xpu.py index 1f44e620..08d2516e 100644 --- a/verl_hardware_plugin/platforms/platform_xpu.py +++ b/verl_hardware_plugin/platforms/platform_xpu.py @@ -77,6 +77,22 @@ class PlatformXPU(PlatformBase): force_sum_reduction to work around this limitation. """ + def __init__(self) -> None: + # Apply this platform's monkeypatch only now, when XPU is actually + # the platform verl selected for this process (verl.plugin.platform. + # platform_manager._create_platform() only constructs PlatformXPU() + # when platform detection -- explicit VERL_PLATFORM=intel or + # auto-detection -- picked "intel"). Applying it from the shared + # verl_hardware_plugin/__init__.py instead would fire on any host + # where XPU hardware/SDK merely happens to be present, even if a + # different platform ends up selected for this run -- e.g. a mixed + # host with VERL_PLATFORM=nvidia explicitly set would still get its + # process-wide torch.distributed.all_reduce(op=AVG) semantics + # changed for no reason. + from verl_hardware_plugin.patches.xpu import reduce_avg_allreduce_patch + + reduce_avg_allreduce_patch.apply() + # ------------------------------------------------------------------ # Core device management # ------------------------------------------------------------------