fix(ppo): warn instead of raising when ref_per_token_logps and ref_input are both passed (#1506)

## Summary

Root cause: in `LigerFusedLinearPPOBase.forward`, the case where both
`ref_per_token_logps` and `ref_input` are passed is handled with `raise
Warning("Both ref_per_token_logps and ref_input are provided. Using
ref_per_token_logps.")`. `Warning` is an exception class, so this line
raises and the GRPO loss call aborts. The rest of the function already
does what the message says: `_ref_input_chunks` is only built when
`ref_per_token_logps is None`, and the reference forward only runs when
`ref_per_token_logps_chunk is None`. So the fallback exists but is never
reached.

Fix: use `warnings.warn(...)` so the call continues and uses
`ref_per_token_logps`, as the message describes.

## Testing Done

Added `test_ref_per_token_logps_preferred_over_ref_input` in
`test/chunked_loss/test_grpo_loss.py`. It passes both
`ref_per_token_logps` and `ref_input`/`ref_weight`, expects a
`UserWarning`, and checks that the loss equals the loss computed with
`ref_per_token_logps` alone.

- Before the fix it fails with `Warning: Both ref_per_token_logps and
ref_input are provided. Using ref_per_token_logps.`
- After the fix it passes, together with the non-compiled
`test_grpo_frozen_lm_head` cases (9 passed).

- Hardware Type: NVIDIA GeForce GTX 1060 6GB, PyTorch 2.5.1 (eager only;
torch.compile needs a newer GPU)
- [x] `make checkstyle` checks: `ruff check .` and `ruff format --check
.` pass (ruff 0.14.11, as pinned in `.pre-commit-config.yaml`)
- [ ] `make test` (not run; Triton is not available on this Windows
host)
- [ ] `make test-convergence` (not run; this change only affects
argument validation)

Co-authored-by: Vaibhav Jindal <[email protected]>

src/liger_kernel/chunked_loss/fused_linear_ppo.py

@@ -1,3 +1,5 @@
+import warnings
+
 from abc import abstractmethod
 from functools import partial
 
@@ -228,7 +230,7 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
                 "If use_ref_model is True, ref_per_token_logps or ref_input must be provided"
             )
             if ref_per_token_logps is not None and ref_input is not None:
-                raise Warning("Both ref_per_token_logps and ref_input are provided. Using ref_per_token_logps.")
+                warnings.warn("Both ref_per_token_logps and ref_input are provided. Using ref_per_token_logps.")
         if loss_type == "dr_grpo":
             assert max_completion_length is not None, "max_completion_length must be provided for loss_type 'dr_grpo'"
         if vllm_is_ratio is not None:

test/chunked_loss/test_grpo_loss.py

@@ -1322,3 +1322,21 @@ def test_selective_logprob_backward_skips_frozen_params(monkeypatch, bias):
         out.sum().backward()
         assert calls[-1] == (weight_requires_grad, bias and bias_requires_grad)
         assert (weight.grad is not None) == weight_requires_grad
+
+
+def test_ref_per_token_logps_preferred_over_ref_input():
+    """Passing both ref_per_token_logps and ref_input warns and uses ref_per_token_logps instead of raising."""
+    B, T, H, V = 2, 8, 16, 32
+    _input = torch.randn(B, T, H, device=device)
+    weight = torch.randn(V, H, device=device)
+    selected_token_ids = torch.randint(0, V, (B, T), device=device)
+    attention_mask = torch.ones(B, T, device=device)
+    advantages = torch.randn(B, device=device)
+    ref_per_token_logps = torch.randn(B, T, device=device) - 3.0
+    loss_fn = LigerFusedLinearGRPOLoss(beta=0.04, use_ref_model=True, compiled=False)
+    args = (_input, weight, selected_token_ids, attention_mask, advantages)
+
+    expected, _ = loss_fn(*args, ref_per_token_logps=ref_per_token_logps)
+    with pytest.warns(UserWarning, match="Using ref_per_token_logps"):
+        loss, _ = loss_fn(*args, ref_per_token_logps=ref_per_token_logps, ref_input=_input, ref_weight=weight)
+    assert_verbose_allclose(loss, expected)