[ops] fix: keep verl FusedLinearForPPO batch-invariant under Liger

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

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"])

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()