LigerFusedLinearGRPOLoss produces ~100x larger grad_norm than TRL's non-Liger-Kernel path due to missing vLLM IS correction and other differences

[yukiu00]

Summary

When using use_liger_kernel=True with TRL v0.27.2's GRPOTrainer and vLLM, grad_norm is ~100x larger than the non-Liger-Kernel path. This is primarily caused by the vLLM importance sampling correction not being applied in the Liger-Kernel loss path. Several other silent differences also exist.

List of differences (TRL v0.27.2)

#	Difference	TRL (non-Liger-Kernel)	Liger-Kernel	Impact	Silent?
1	vLLM IS correction	per_token_loss *= importance_sampling_ratio (L2351-L2352)	Not applied	~100-300x (primary cause)	Yes
2	dapo/cispo normalizer	num_items_in_batch / num_processes (total tokens across entire generation batch) (L2371)	all_reduce(sum(attention_mask)) / world_size (current micro-batch only), then divided by current_gradient_accumulation_steps (Liger-Kernel L116-117, TRL L2194)	Several-fold depending on length variance	Yes
3	tool_mask	completion_mask * tool_mask (L2243)	completion_mask only (TRL L2178)	Proportional to tool token ratio	Yes
4	use_bias_correction_kl	per_token_kl *= coef_1 (L2315-L2316)	Not implemented	KL term only	Yes
5	delta (ratio clamping)	coef_1 = clamp(coef_1, max=delta) (L2326-L2327)	Not implemented	Changes clipping behavior	Yes
6	off_policy_mask, top_entropy_quantile, sequence-level IS, sapo	Supported	Not supported	-	No (raises error)

Where to fix

#	Difference	Where to fix	Notes
1	vLLM IS correction	Liger-Kernel + TRL	Liger-Kernel: add parameter to apply per_token_loss *= is_ratio before reduction. TRL: pass it from compute_liger_loss
2	dapo/cispo normalizer	Liger-Kernel + TRL	Liger-Kernel: accept an external normalizer. TRL: pass num_items_in_batch
3	tool_mask	TRL only	Pass completion_mask * tool_mask in compute_liger_loss
4	use_bias_correction_kl	Liger-Kernel + TRL	Liger-Kernel: add flag to ppo_loss_fn. TRL: pass it
5	delta (ratio clamping)	Liger-Kernel + TRL	Liger-Kernel: add clamp argument for coef_1. TRL: pass it

Observed behavior

With DAPO loss on a large model using vLLM in GRPOTrainer:

Metric (step 1)	Liger-Kernel (use_liger_kernel=True)	Non-Liger-Kernel
grad_norm	31.09	0.29
loss	0.0187	-0.0001

The IS ratio mean was ~0.003–0.01, so the non-Liger-Kernel path scales down the loss by that factor, while the Liger-Kernel path passes it through as-is.

Conclusion

The primary cause is # 1: vLLM IS correction is not applied in Liger-Kernel. Since the IS ratio must be multiplied before reduction, it cannot be fixed on the TRL side alone — a Liger-Kernel change is necessary. Only # 3 can be resolved with a one-line fix on the TRL side.

[yukiu00]

Minimal repro (vLLM IS correction)

I used a direct GRPOTrainer.compute_loss() repro (no rollout/generation), with fixed tensors and 5 steps.

import torch
from trl import GRPOConfig, GRPOTrainer

MODEL = "HuggingFaceTB/SmolLM2-135M"
B, P, T = 2, 8, 16
IS_RATIO = 0.01
STEPS = 5
BASE_SEED = 42

def reward_fn(prompts, completions, **kwargs):
    return [1.0] * len(prompts)

def make_inputs(model, step):
    device = next(model.parameters()).device
    vocab = model.config.vocab_size
    g = torch.Generator().manual_seed(BASE_SEED + step)
    prompt_ids = torch.randint(10, vocab, (B, P), generator=g, dtype=torch.long).to(device)
    completion_ids = torch.randint(10, vocab, (B, T), generator=g, dtype=torch.long).to(device)
    prompt_mask = torch.ones_like(prompt_ids)
    completion_mask = torch.ones_like(completion_ids)
    return {
        "prompt_ids": prompt_ids,
        "prompt_mask": prompt_mask,
        "completion_ids": completion_ids,
        "completion_mask": completion_mask,
        "advantages": torch.ones(B, device=device),
        "old_per_token_logps": torch.full((B, T), -2.0, device=device),
        "importance_sampling_ratio": torch.full((B, T), IS_RATIO, device=device),
        "num_items_in_batch": float(completion_mask.sum().item()),
    }

def grad_norm(model):
    s = torch.zeros((), device=next(model.parameters()).device)
    for p in model.parameters():
        if p.grad is not None:
            s = s + p.grad.float().pow(2).sum()
    return s.sqrt().item()

def build_case(use_liger, use_vllm_is):
    args = GRPOConfig(
        output_dir="tmp_grpo_issue_repro",
        max_steps=1,
        per_device_train_batch_size=2,
        num_generations=2,
        gradient_accumulation_steps=1,
        use_liger_kernel=use_liger,
        use_vllm=False,  # avoid server init in repro
        vllm_importance_sampling_correction=use_vllm_is,
        beta=0.0,
        loss_type="dapo",
        importance_sampling_level="token",
        bf16=False,
        fp16=False,
        report_to="none",
        save_strategy="no",
        disable_tqdm=True,
    )
    trainer = GRPOTrainer(model=MODEL, reward_funcs=reward_fn, args=args, train_dataset=None)
    trainer.use_vllm = True
    trainer.vllm_importance_sampling_correction = use_vllm_is
    trainer.current_gradient_accumulation_steps = trainer.args.gradient_accumulation_steps
    model = trainer.model.cuda()
    model.train()
    return trainer, model

def run_step(trainer, model, step):
    model.zero_grad(set_to_none=True)
    loss = trainer.compute_loss(model, make_inputs(model, step))
    loss.backward()
    return grad_norm(model)

def run_case_steps(use_liger, use_vllm_is):
    trainer, model = build_case(use_liger, use_vllm_is)
    return [run_step(trainer, model, s) for s in range(STEPS)]

if __name__ == "__main__":
    non_off = run_case_steps(False, False)
    non_on = run_case_steps(False, True)
    liger_off = run_case_steps(True, False)
    liger_on = run_case_steps(True, True)

    print("step | non_liger(off/on ratio) | liger(off/on ratio)")
    for s in range(STEPS):
        ratio_non = non_on[s] / non_off[s]
        ratio_liger = liger_on[s] / liger_off[s]
        print(f"{s+1:>4} | non_liger: {non_off[s]:.6e}->{non_on[s]:.6e}, ratio={ratio_non:.6f} | "
              f"liger: {liger_off[s]:.6e}->{liger_on[s]:.6e}, ratio={ratio_liger:.6f}")

Observed output

step | non_liger(off/on ratio) | liger(off/on ratio)
   1 | non_liger: 9.116095e-03->9.116097e-05, ratio=0.010000 | liger: 9.116122e-03->9.116122e-03, ratio=1.000000
   2 | non_liger: 8.003094e-03->8.003095e-05, ratio=0.010000 | liger: 8.003096e-03->8.003096e-03, ratio=1.000000
   3 | non_liger: 1.983273e-03->1.983274e-05, ratio=0.010000 | liger: 1.983265e-03->1.983265e-03, ratio=1.000000
   4 | non_liger: 2.166789e-03->2.166788e-05, ratio=0.010000 | liger: 2.166775e-03->2.166775e-03, ratio=1.000000
   5 | non_liger: 1.514107e-03->1.514108e-05, ratio=0.010000 | liger: 1.514113e-03->1.514113e-03, ratio=1.000000

mean ratio(non-liger, IS on/off): 0.010000
mean ratio(liger, IS on/off): 1.000000

[kashif]

@yukiu00 can this issue be closed?