[training_utils, env, doc] feat: use Liger fused linear PPO kernel

docs/examples/config.rst

@@ -749,4 +749,4 @@ Most parameters for Model are similar to Reward Model.
   default to ``all-linear``. See `peft docs <https://huggingface.co/docs/peft/v0.15.0/en/package_reference/lora#peft.LoraConfig.target_modules>`_ for detail.
 
 - ``use_liger``: Whether to enable Liger kernel, default to False. If True,
-  we apply Liger kernel to the model (depends on `liger-kernel`).
+  we apply Liger kernel to the model (depends on ``liger-kernel>=0.8.2``).

docs/perf/perf_tuning.rst

@@ -174,16 +174,16 @@ LigerKernel for training performance
 
 LigerKernel provides fused Triton kernels (RMSNorm, SwiGLU, RoPE) that can improve training throughput. It works with both SFT and RL (PPO/GRPO) training, including vision-language models.
 
-1. Install liger-kernel via ``pip3 install liger-kernel``. Set ``use_liger`` in your configuration:
+1. Install Liger Kernel 0.8.2 or newer via ``pip3 install "liger-kernel>=0.8.2"``. Set ``use_liger`` in your configuration:
 
    .. code-block:: yaml
 
       model:
         use_liger: True  # Enable LigerKernel
 
-2. The default value is ``False``. When enabled, verl applies Liger's fused RMSNorm, SwiGLU, and RoPE kernels to the model. The ``fused_linear_cross_entropy`` optimization is disabled because verl computes log-probabilities via its own path.
+2. The default value is ``False``. When enabled, verl applies Liger's fused RMSNorm, SwiGLU, and RoPE kernels to the model. The model-level ``fused_linear_cross_entropy`` patch remains disabled because verl computes log-probabilities through its output-head path. With ``use_fused_kernels`` and the ``torch`` backend, that path uses Liger's fused scaled linear cross entropy from v0.8.2 or newer and falls back to verl's existing chunked ``FusedLinearForPPOFunction`` when Liger is not installed.
 
-3. ``use_liger`` is compatible with ``use_fused_kernels`` — they operate at different levels (Liger optimizes model internals, fused kernels optimize the output head). Using both together gives the best speed-memory tradeoff.
+3. ``use_liger`` is compatible with ``use_fused_kernels``. The former controls model-internal kernels, while the latter controls the output head and can use Liger's scaled cross entropy independently when ``liger-kernel>=0.8.2`` is installed.
 
 Forward prefetch in FSDP training backend
 ----------------------

docs/start/install.rst

@@ -405,7 +405,7 @@ Find the docker for AMD ROCm: `docker/Dockerfile.rocm <https://github.com/verl-p
         datasets \
         dill \
         hydra-core \
-        liger-kernel \
+        "liger-kernel>=0.8.2" \
         numpy \
         pandas \
         datasets \

pyproject.toml

@@ -140,7 +140,7 @@ fsdp = [
     "torch==2.11.0",
     "torchvision==0.26.0",
     "torchaudio==2.11.0",
-    "liger-kernel",
+    "liger-kernel>=0.8.2",
     "trl==0.27.0",
     # transformers NOT pinned here: trainers inherit the synced engine's version
     # (vllm 5.5.3 / sglang 5.3.0). See override-dependencies below.
@@ -162,7 +162,7 @@ megatron = [
     "onnxscript",
     "trl==0.27.0",
     "matplotlib",
-    "liger-kernel",
+    "liger-kernel>=0.8.2",
     # Two Megatron<->HF connectors: mbridge (legacy) and megatron-bridge (NVIDIA,
     # imported as megatron.bridge; the migration target, installed deps-free — see
     # [[tool.uv.dependency-metadata]] below). 0.5.2 is the wheelhouse build of the

requirements.txt

@@ -4,7 +4,7 @@ codetiming
 datasets
 dill
 hydra-core
-liger-kernel
+liger-kernel>=0.8.2
 numpy>=2.0.0
 pandas
 peft

setup.py

@@ -50,7 +50,7 @@
 # Empty since PRIME code scoring dropped `pyext`; kept so `verl[prime]` stays installable.
 PRIME_REQUIRES = []
 GEO_REQUIRES = ["mathruler", "torchvision", "qwen_vl_utils"]
-GPU_REQUIRES = ["liger-kernel", "flash-attn"]
+GPU_REQUIRES = ["liger-kernel>=0.8.2", "flash-attn"]
 MATH_REQUIRES = ["math-verify"]  # Add math-verify as an optional dependency
 VLLM_REQUIRES = ["tensordict>=0.8.0,<=0.10.0,!=0.9.0", "vllm>=0.18.0"]
 TRTLLM_REQUIRES = ["tensorrt-llm>=1.2.0rc6"]

tests/utils/test_experimental_torch_functional_on_cpu.py

@@ -0,0 +1,117 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+import torch
+
+_MODULE_PATH = Path(__file__).resolve().parents[2] / "verl" / "utils" / "experimental" / "torch_functional.py"
+_SPEC = importlib.util.spec_from_file_location("_experimental_torch_functional", _MODULE_PATH)
+assert _SPEC is not None and _SPEC.loader is not None
+experimental_F = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(experimental_F)
+
+
+@pytest.mark.parametrize("hidden_shape", [(7, 5), (2, 7, 5)])
+def test_fused_linear_for_ppo_chunked_fallback_matches_torch(monkeypatch, hidden_shape):
+    monkeypatch.setattr(experimental_F, "_LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY", None)
+    monkeypatch.setattr(experimental_F, "_FLASH_ATTN_CROSS_ENTROPY_AVAILABLE", False)
+    torch.manual_seed(42)
+
+    temperature = 0.7
+    vocab_size = 11
+    hidden = torch.randn(hidden_shape, requires_grad=True)
+    weight = torch.randn(vocab_size, hidden_shape[-1], requires_grad=True)
+    labels = torch.randint(vocab_size, hidden_shape[:-1], dtype=torch.int32)
+    grad_log_probs = torch.randn(hidden_shape[:-1])
+    grad_entropy = torch.randn(hidden_shape[:-1])
+
+    log_probs, entropy = experimental_F.FusedLinearForPPO()(hidden, weight, labels, temperature)
+    torch.autograd.backward((log_probs, entropy), (grad_log_probs, grad_entropy))
+
+    expected_hidden = hidden.detach().clone().requires_grad_(True)
+    expected_weight = weight.detach().clone().requires_grad_(True)
+    logits = ((expected_hidden @ expected_weight.t()) / temperature).float()
+    expected_log_probs = logits.log_softmax(dim=-1).gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)
+    probs = logits.softmax(dim=-1)
+    expected_entropy = torch.logsumexp(logits, dim=-1) - torch.sum(probs * logits, dim=-1)
+    torch.autograd.backward(
+        (expected_log_probs, expected_entropy),
+        (grad_log_probs, grad_entropy),
+    )
+
+    torch.testing.assert_close(log_probs, expected_log_probs)
+    torch.testing.assert_close(entropy, expected_entropy)
+    torch.testing.assert_close(hidden.grad, expected_hidden.grad)
+    torch.testing.assert_close(weight.grad, expected_weight.grad)
+
+
+def test_fused_linear_for_ppo_dispatches_to_liger(monkeypatch):
+    calls = []
+
+    class FakeLigerFusedLinearScaledCrossEntropyFunction:
+        @staticmethod
+        def apply(*args):
+            calls.append(args)
+            hidden_states = args[0]
+            token_count = hidden_states.shape[0]
+            nll = torch.arange(token_count, dtype=torch.float32)
+            entropy = torch.arange(token_count, dtype=hidden_states.dtype) + 10
+            return nll, entropy
+
+    monkeypatch.setattr(
+        experimental_F,
+        "_LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY",
+        FakeLigerFusedLinearScaledCrossEntropyFunction,
+    )
+    hidden = torch.randn(2, 3, 5)
+    weight = torch.randn(7, 5)
+    labels = torch.randint(7, (2, 3), dtype=torch.int32)
+
+    log_probs, entropy = experimental_F.FusedLinearForPPO()(hidden, weight, labels, temperature=0.8)
+
+    assert len(calls) == 1
+    liger_hidden, liger_weight, liger_labels, temperature, ignore_index, m_tiles, return_entropy = calls[0]
+    assert liger_hidden.shape == (6, 5)
+    assert liger_weight is weight
+    assert liger_labels.shape == (6,)
+    assert liger_labels.dtype == torch.int64
+    assert temperature == 0.8
+    assert ignore_index == -100
+    assert m_tiles == 1
+    assert return_entropy is True
+    torch.testing.assert_close(log_probs, -torch.arange(6, dtype=torch.float32).reshape(2, 3))
+    torch.testing.assert_close(entropy, (torch.arange(6, dtype=hidden.dtype) + 10).reshape(2, 3))
+
+
+def test_fused_linear_for_ppo_fallback_preserves_chunking(monkeypatch):
+    monkeypatch.setattr(experimental_F, "_LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY", None)
+    monkeypatch.setattr(experimental_F, "_FLASH_ATTN_CROSS_ENTROPY_AVAILABLE", False)
+    original_forward = experimental_F._fused_linear_for_ppo_fwd
+    chunk_sizes = []
+
+    def record_chunk_size(hidden_states, *args, **kwargs):
+        chunk_sizes.append(hidden_states.shape[0])
+        return original_forward(hidden_states, *args, **kwargs)
+
+    monkeypatch.setattr(experimental_F, "_fused_linear_for_ppo_fwd", record_chunk_size)
+    hidden = torch.randn(1, 7, 5)
+    weight = torch.randn(11, 5)
+    labels = torch.randint(11, (1, 7))
+
+    experimental_F.FusedLinearForPPO(chunk_size=3)(hidden, weight, labels)
+
+    assert chunk_sizes == [3, 3, 1]

verl/utils/experimental/torch_functional.py

@@ -16,6 +16,15 @@
 
 import torch
 
+try:
+    from liger_kernel.ops import (
+        LigerFusedLinearScaledCrossEntropyFunction as _LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY,
+    )
+except ModuleNotFoundError as exc:
+    if exc.name != "liger_kernel":
+        raise
+    _LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY = None
+
 try:
     from flash_attn.ops.triton.cross_entropy import cross_entropy_loss
 
@@ -222,10 +231,36 @@ def forward(
         temperature: float = 1.0,
     ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
         input_ids = input_ids.to(torch.int64)
-        return FusedLinearForPPOFunction.apply(
+        if _LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY is None:
+            return FusedLinearForPPOFunction.apply(
+                hidden_states,
+                vocab_weights,
+                input_ids,
+                temperature,
+                self.chunk_size,
+            )
+
+        if hidden_states.ndim not in (2, 3):
+            raise ValueError(f"hidden_states must be 2D or 3D, got shape {tuple(hidden_states.shape)}")
+        if input_ids.shape != hidden_states.shape[:-1]:
+            raise ValueError(
+                f"input_ids shape {tuple(input_ids.shape)} must match hidden_states shape "
+                f"{tuple(hidden_states.shape[:-1])}"
+            )
+
+        output_shape = input_ids.shape
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+        input_ids = input_ids.reshape(-1)
+
+        nll, entropy = _LIGER_FUSED_LINEAR_SCALED_CROSS_ENTROPY.apply(
             hidden_states,
             vocab_weights,
             input_ids,
             temperature,
-            self.chunk_size,
+            -100,
+            1,
+            True,
         )
+        log_probs = -nll
+
+        return log_probs.reshape(output_shape), entropy.reshape(output_shape)