Add vLLM importance sampling ratio support for GRPO loss

[yukiu00]

Summary

Fixes the primary cause (item 1) of #1082 — LigerFusedLinearGRPOLoss produces ~100x larger grad_norm than TRL's non-Liger path when using vLLM.

Root cause: TRL's GRPOTrainer applies per_token_loss *= importance_sampling_ratio (source) to correct for distribution mismatch from vLLM's rejection/stratified sampling. Liger-Kernel had no mechanism to accept or apply this correction, so the IS ratio was silently ignored, resulting in uncorrected (and much larger) gradients.

This is a high-priority fix — any user running GRPOTrainer with use_vllm=True and use_liger_kernel=True is affected, and the resulting ~100x gradient mismatch can cause training instability or divergence.

Changes

• Add optional vllm_is_ratio parameter ([B, T] tensor or None) to both code paths:
  • Chunked loss path: LigerFusedLinearGRPOLoss, LigerFusedLinearGRPOFunction, ppo_loss_fn, and the base class LigerFusedLinearPPOBase chunking pipeline
  • Triton kernel path: triton_grpo_loss, GrpoLossFunction, and the Triton fwd/bwd kernels (_grpo_loss_fwd_kernel, _grpo_loss_bwd_kernel)
• The IS correction is applied after PPO clipped loss computation and before KL penalty, matching TRL's behavior exactly
• vllm_is_ratio=None (default) preserves existing behavior — no breaking changes
• Works with all loss types: grpo, dapo, bnpo, dr_grpo, cispo, sapo

Verification

With IS_RATIO=0.01, the grad_norm ratio matches exactly:

Chunked loss path:
  grad_norm WITHOUT vllm_is_ratio: 1.052219e-01
  grad_norm WITH    vllm_is_ratio: 1.052219e-03
  ratio: 0.010000 ✓

Triton path:
  grad_norm WITHOUT vllm_is_ratio: 1.461673e-02
  grad_norm WITH    vllm_is_ratio: 1.461673e-04
  ratio: 0.010000 ✓

Test plan

• Extended existing test_correctness in test/chunked_loss/test_grpo_loss.py with use_vllm_is_ratio parametrize — covers all 6 loss types × 2 IS levels × 2 beta values × with/without vllm_is_ratio
• Added test_grpo_loss_with_vllm_is_ratio in test/transformers/test_grpo_loss.py — compares Triton output against PyTorch reference with IS correction, plus vllm_is_ratio=None == vllm_is_ratio=ones identity check
• All existing tests continue to pass (no regressions)
• make checkstyle passes

Related

• Reference implementation: [GRPO] add grpo loss types #993
• Issue: LigerFusedLinearGRPOLoss produces ~100x larger grad_norm than TRL's non-Liger-Kernel path due to missing vLLM IS correction and other differences #1082

[yukiu00]

@Tcc0403 Could you please review this PR? This fixes a critical issue (#1082) where LigerFusedLinearGRPOLoss produces ~100x larger grad_norm than TRL's non-Liger path when using vLLM. Since many users rely on GRPOTrainer with vLLM + Liger in production, this bug can silently cause training instability. An early review and merge would help unblock affected users.

[kashif]

thanks!

[Tcc0403]

LGTM, thanks for these high quality PRs!

[yukiu00]

@Tcc0403 Thanks for the quick review and merge!

Could you let me know which version this change will be released in? I have a downstream PR in TRL (huggingface/trl#5031) that depends on this, and I need to set the correct minimum version for the version check.

[Tcc0403]

@vaibhavjindal @Mecoli1219 @momochen Can we have a release for these fixes?