diff --git a/tests/batch_invariant_ops/test_verl_fused_linear_for_ppo.py b/tests/batch_invariant_ops/test_verl_fused_linear_for_ppo.py index aeb2715..b820786 100644 --- a/tests/batch_invariant_ops/test_verl_fused_linear_for_ppo.py +++ b/tests/batch_invariant_ops/test_verl_fused_linear_for_ppo.py @@ -17,6 +17,7 @@ import pytest import torch +import verl.utils.experimental.torch_functional as verl_torch_functional from verl.utils.experimental.torch_functional import FusedLinearForPPO from verl.utils.torch_functional import logprobs_from_logits from vexact.batch_invariant_ops import set_batch_invariant_mode @@ -132,3 +133,24 @@ def test_fused_linear_for_ppo_log_probs_float32_in_bfloat16_model( rtol=0, atol=0, ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="FusedLinearForPPO requires CUDA") +@pytest.mark.skipif( + not hasattr(verl_torch_functional, "_LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY"), + reason="this verl has no Liger dispatch for FusedLinearForPPO", +) +def test_batch_invariant_mode_disables_verl_liger_dispatch() -> None: + """ + checks that batch-invariant mode sends verl FusedLinearForPPO down its chunked + torch.matmul path (Liger dispatch nulled) and restores the dispatch afterwards + """ + original = verl_torch_functional._LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY + + with set_batch_invariant_mode(True): + assert verl_torch_functional._LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY is None + with set_batch_invariant_mode(False): + assert verl_torch_functional._LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY is original + assert verl_torch_functional._LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY is None + + assert verl_torch_functional._LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY is original diff --git a/vexact/batch_invariant_ops/batch_invariant_ops.py b/vexact/batch_invariant_ops/batch_invariant_ops.py index da077f4..9d47a7a 100644 --- a/vexact/batch_invariant_ops/batch_invariant_ops.py +++ b/vexact/batch_invariant_ops/batch_invariant_ops.py @@ -7,6 +7,7 @@ import contextlib import os +import sys from collections import namedtuple from collections.abc import Callable from typing import Any @@ -738,9 +739,46 @@ def is_batch_invariant_mode_enabled(): return _batch_invariant_MODE +# verl's ``FusedLinearForPPO`` (verl-project/verl#7461) routes the fused output head +# through Liger's ``LigerFusedLinearScaledCrossEntropyFunction`` whenever ``liger_kernel`` +# is importable. That kernel bypasses the ``aten::mm`` override below, so the training-side +# log-probs no longer match vexact's inference log-probs bit-for-bit. While batch-invariant +# mode is on we null verl's module-level dispatch handle, which sends ``FusedLinearForPPO`` +# back to its chunked ``torch.matmul`` + cross-entropy path (the one the overrides cover), +# and restore it when the mode is turned off. +_VERL_TORCH_FUNCTIONAL_MODULE = "verl.utils.experimental.torch_functional" +_VERL_LIGER_DISPATCH_ATTR = "_LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY" +# (module, original dispatch handle) while the hook is applied, else None. +_verl_liger_dispatch_backup: tuple[Any, Any] | None = None + + +def _disable_verl_liger_fused_linear() -> None: + """Null verl's Liger dispatch handle if verl's fused-linear module is loaded.""" + global _verl_liger_dispatch_backup + if _verl_liger_dispatch_backup is not None: + return + module = sys.modules.get(_VERL_TORCH_FUNCTIONAL_MODULE) + if module is None or not hasattr(module, _VERL_LIGER_DISPATCH_ATTR): + return + _verl_liger_dispatch_backup = (module, getattr(module, _VERL_LIGER_DISPATCH_ATTR)) + setattr(module, _VERL_LIGER_DISPATCH_ATTR, None) + + +def _restore_verl_liger_fused_linear() -> None: + """Undo :func:`_disable_verl_liger_fused_linear`.""" + global _verl_liger_dispatch_backup + if _verl_liger_dispatch_backup is None: + return + module, original = _verl_liger_dispatch_backup + setattr(module, _VERL_LIGER_DISPATCH_ATTR, original) + _verl_liger_dispatch_backup = None + + def enable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB if _batch_invariant_MODE: + # verl's module may have been imported after the mode was enabled; re-apply the hook. + _disable_verl_liger_fused_linear() return dispatch_key = getattr(torch.accelerator.current_accelerator(), "type", "cpu").upper() @@ -757,6 +795,7 @@ def enable_batch_invariant_mode(): torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False torch.use_deterministic_algorithms(True, warn_only=False) torch.backends.cudnn.deterministic = True + _disable_verl_liger_fused_linear() def disable_batch_invariant_mode(): @@ -765,6 +804,7 @@ def disable_batch_invariant_mode(): _batch_invariant_LIB._destroy() _batch_invariant_MODE = False _batch_invariant_LIB = None + _restore_verl_liger_fused_linear() @contextlib.contextmanager @@ -779,6 +819,10 @@ def set_batch_invariant_mode(enabled: bool = True): if _batch_invariant_LIB is not None: _batch_invariant_LIB._destroy() _batch_invariant_MODE, _batch_invariant_LIB = old_data + if _batch_invariant_MODE: + _disable_verl_liger_fused_linear() + else: + _restore_verl_liger_fused_linear() AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"]) diff --git a/vexact/integrations/verl/fsdp_enable_invariant.py b/vexact/integrations/verl/fsdp_enable_invariant.py index c8fe012..7347c4e 100644 --- a/vexact/integrations/verl/fsdp_enable_invariant.py +++ b/vexact/integrations/verl/fsdp_enable_invariant.py @@ -14,16 +14,18 @@ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +# Load verl's fused-linear module before enabling batch-invariant mode so the mode can +# null its Liger dispatch handle (see ``_disable_verl_liger_fused_linear``). +import verl.utils.experimental.torch_functional # noqa: F401 from vexact.batch_invariant_ops import ( enable_batch_invariant_mode, - is_batch_invariant_mode_enabled, triton_flash_attention_forward, ) from vexact.models.register import register_models as _register_models -if not is_batch_invariant_mode_enabled(): - enable_batch_invariant_mode() +# Also re-applies the verl Liger hook when the mode was already enabled by an earlier import. +enable_batch_invariant_mode() ALL_ATTENTION_FUNCTIONS["triton-invariant"] = triton_flash_attention_forward _register_models()