From 27eb5ce3ab3c56ad90b47e1a4201ae121e09b764 Mon Sep 17 00:00:00 2001 From: chinoll Date: Thu, 23 Jul 2026 04:10:36 +0800 Subject: [PATCH 1/3] Add GMPO to fused linear GRPO loss --- src/liger_kernel/chunked_loss/grpo_loss.py | 52 +-- test/chunked_loss/test_grpo_loss.py | 347 +++++++++++++++++++++ 2 files changed, 381 insertions(+), 18 deletions(-) diff --git a/src/liger_kernel/chunked_loss/grpo_loss.py b/src/liger_kernel/chunked_loss/grpo_loss.py index d6bbd0a14..087f981a5 100644 --- a/src/liger_kernel/chunked_loss/grpo_loss.py +++ b/src/liger_kernel/chunked_loss/grpo_loss.py @@ -74,8 +74,17 @@ def sapo_loss_fn(importance_ratio: torch.Tensor, temperature: float) -> torch.Te return sigmoid_smoothed_loss * 4 / temperature -def clip_coef_fn(coef, epsilon_low, epsilon_high, loss_type): - if loss_type == "cispo": +def clip_coef_fn(coef, epsilon_low, epsilon_high, loss_type, advantages=None): + if loss_type == "gmpo": + positive_advantages = advantages.unsqueeze(1) >= 0 + clipped_coef = torch.where( + positive_advantages, + torch.minimum(coef, torch.as_tensor(epsilon_high, device=coef.device, dtype=coef.dtype)), + torch.maximum(coef, torch.as_tensor(-epsilon_low, device=coef.device, dtype=coef.dtype)), + ) + is_lower_clipped = coef < -epsilon_low + is_upper_clipped = coef > epsilon_high + elif loss_type == "cispo": # CISPO: clip and detach the importance weights upper_bound = epsilon_high lower_bound = None @@ -108,7 +117,7 @@ def ppo_loss_fn( epsilon_low=0.2, epsilon_high=0.2, beta=0.04, - loss_type="dapo", # ["grpo", "bnpo", "dr_grpo", "dapo", "cispo", "sapo", "luspo", "vespo"] + loss_type="dapo", # ["grpo", "gmpo", "bnpo", "dr_grpo", "dapo", "cispo", "sapo", "luspo", "vespo"] max_completion_length=None, # Required for dr_grpo importance_sampling_level="token", # ["token", "sequence"] - new parameter for GSPO sapo_temperature_pos=1.0, # Temperature for positive advantages in SAPO @@ -125,7 +134,7 @@ def ppo_loss_fn( ): """GRPO Loss Function matching GRPOTrainer implementation.""" # Validate sequence-level + loss_type combinations - if importance_sampling_level == "sequence" and loss_type in ("cispo", "sapo", "vespo"): + if importance_sampling_level == "sequence" and loss_type in ("gmpo", "cispo", "sapo", "vespo"): raise ValueError( f"Sequence-level importance sampling is not supported for loss_type='{loss_type}'. " f"Use importance_sampling_level='token' instead." @@ -150,11 +159,18 @@ def ppo_loss_fn( "and 'sequence'." ) - # From here, log_importance_weights (and all subsequent tensors, coef_1, coef_2, etc.) shape depends on - # importance_sampling_level: "token" level: (B, T); "sequence" level: (B, 1) - coef_1 = torch.exp(log_importance_weights) - coef_2, is_lower_clipped, is_upper_clipped = clip_coef_fn(coef_1, epsilon_low, epsilon_high, loss_type) - if loss_type == "cispo": + # GMPO coefficients are log-domain; all other loss types use ratio-domain coefficients. + coef_1 = log_importance_weights if loss_type == "gmpo" else torch.exp(log_importance_weights) + coef_2, is_lower_clipped, is_upper_clipped = clip_coef_fn( + coef_1, epsilon_low, epsilon_high, loss_type, advantages=advantages + ) + + if loss_type == "gmpo": + token_counts = attention_mask.sum(-1).clamp(min=1.0) + mean_log_ratio = (coef_2 * attention_mask).sum(-1) / token_counts + sample_ratio = torch.exp(mean_log_ratio) + per_token_loss = (-advantages.unsqueeze(1) * sample_ratio.unsqueeze(1)).expand_as(attention_mask) + elif loss_type == "cispo": # CISPO: clip and detach the importance weights, multiply by log probs # Reference: https://github.com/huggingface/trl/blob/035c3ff151b953ca72cdfe0ee966bc1469a26fde/trl/trainer/grpo_trainer.py#L2030 per_token_loss = -coef_2 * advantages.unsqueeze(1) * per_token_logps @@ -202,27 +218,26 @@ def ppo_loss_fn( # Apply vLLM importance sampling correction BEFORE adding KL penalty # VESPO folds this correction into phi_seq (in log space), so we skip it here. - if vllm_is_ratio is not None and loss_type != "vespo": + if vllm_is_ratio is not None and loss_type not in ("gmpo", "vespo"): per_token_loss = per_token_loss * vllm_is_ratio if beta != 0.0: # Compute KL penalty (approximates KL[per_token_logps, ref_per_token_logps]) kl_div = k3_loss_fn(ref_per_token_logps, per_token_logps) if use_bias_correction_kl: - # Importance-sampling-corrected KL (DeepSeek-V3.2): kl *= coef_1. + # Importance-sampling-corrected KL (DeepSeek-V3.2): kl *= unclipped importance ratio. # Use exp(log_importance_weights) so the ratio's shape matches # importance_sampling_level (token: (B, T); sequence: (B, 1)), # mirroring TRL's ``per_token_kl * coef_1`` (un-clamped, before delta). kl_div = kl_div * torch.exp(log_importance_weights) - # Combine losses per_token_loss = per_token_loss + beta * kl_div # Note: We normalize by the number of tokens in the batch (using full_attention_mask), # which is consistent with the DAPO loss implementation (https://arxiv.org/html/2503.14476v1) # and TRL GRPO implementation # (https://github.com/huggingface/trl/blob/e751a16df56e70190fb94bed4a2035eec3303777/trl/trainer/grpo_trainer.py#L966) - if loss_type == "grpo" or loss_type == "sapo": - # Average per-sequence loss (SAPO uses same normalization as GRPO) + if loss_type in ("grpo", "gmpo", "sapo"): + # Average per-sequence loss (GMPO and SAPO use the same normalization as GRPO) loss = ( (per_token_loss * attention_mask).sum(-1) / torch.clamp(attention_mask.sum(-1), min=1.0) ).sum() / full_attention_mask.shape[0] @@ -250,7 +265,6 @@ def ppo_loss_fn( metrics = [] if beta != 0.0: metrics.append(((kl_div * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0))) - # Adjust clipping metric calculation based on importance sampling level if importance_sampling_level == "token": is_clipped = (is_lower_clipped & (advantages.unsqueeze(1) < 0)) | ( @@ -316,8 +330,9 @@ def forward( ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size) ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,) beta (float): Weight for the KL penalty - loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo", "dapo", "cispo", "sapo", "luspo"). + loss_type (str): Type of loss calculation ("grpo", "gmpo", "bnpo", "dr_grpo", "dapo", "cispo", "sapo", "luspo"). Defaults to "dapo". + GMPO treats epsilon_low/high as log-space bounds ``[-epsilon_low, +epsilon_high]``. max_completion_length (int, optional): Maximum completion length, required for "dr_grpo". Defaults to None. importance_sampling_level (str): Level of importance sampling ("token" or "sequence"). Defaults to "token". sapo_temperature_pos (float): Temperature for positive advantages in SAPO. Defaults to 1.0. @@ -332,7 +347,7 @@ def forward( torch.Tensor: Computed loss """ # Validate before entering torch.compile boundary - if importance_sampling_level == "sequence" and loss_type in ("cispo", "sapo", "vespo"): + if importance_sampling_level == "sequence" and loss_type in ("gmpo", "cispo", "sapo", "vespo"): raise ValueError( f"Sequence-level importance sampling is not supported for loss_type='{loss_type}'. " f"Use importance_sampling_level='token' instead." @@ -447,9 +462,10 @@ def __init__( chunk_size (int): Size of chunks for processing. epsilon_low (float): Lower bound for the importance sampling ratio. epsilon_high (float): Upper bound for the importance sampling ratio. - loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo", "dapo", "cispo", "sapo", "luspo"). + loss_type (str): Type of loss calculation ("grpo", "gmpo", "bnpo", "dr_grpo", "dapo", "cispo", "sapo", "luspo"). Defaults to "dapo". For "cispo", epsilon_high is typically larger (e.g. 5.0) and epsilon_low is unused. For "sapo", uses soft gating instead of hard clipping. + GMPO treats epsilon_low/high as log-space bounds ``[-epsilon_low, +epsilon_high]``. max_completion_length (int, optional): Maximum completion length, required for "dr_grpo". Defaults to None. importance_sampling_level (str): Level of importance sampling ("token" or "sequence"). Defaults to "token". sapo_temperature_pos (float): Temperature for positive advantages in SAPO. Defaults to 1.0. diff --git a/test/chunked_loss/test_grpo_loss.py b/test/chunked_loss/test_grpo_loss.py index 9da0aedf8..e6b238817 100644 --- a/test/chunked_loss/test_grpo_loss.py +++ b/test/chunked_loss/test_grpo_loss.py @@ -31,6 +31,45 @@ def sapo_loss_fn(importance_ratio: torch.Tensor, temperature: float) -> torch.Te return sigmoid_smoothed_loss * 4 / temperature +def torch_gmpo_loss( + per_token_logps, + attention_mask, + advantages, + full_attention_mask, + old_per_token_logps, + ref_per_token_logps, + epsilon_low, + epsilon_high, + beta, + use_bias_correction_kl=False, +): + log_ratio = per_token_logps - old_per_token_logps + lengths = attention_mask.sum(-1).clamp(min=1.0) + positive_advantages = advantages.unsqueeze(1) >= 0 + effective_log_ratio = torch.where( + positive_advantages, + torch.minimum(log_ratio, torch.as_tensor(epsilon_high, device=log_ratio.device, dtype=log_ratio.dtype)), + torch.maximum(log_ratio, torch.as_tensor(-epsilon_low, device=log_ratio.device, dtype=log_ratio.dtype)), + ) + sequence_log_ratio = (effective_log_ratio * attention_mask).sum(-1) / lengths + loss = (-advantages * torch.exp(sequence_log_ratio)).sum() / full_attention_mask.shape[0] + + metrics = [] + if beta != 0.0: + ref_delta = ref_per_token_logps - per_token_logps + kl_div = torch.exp(ref_delta) - ref_delta - 1.0 + if use_bias_correction_kl: + kl_div = kl_div * torch.exp(log_ratio) + per_sequence_kl = (kl_div * attention_mask).sum(-1) / lengths + loss = loss + beta * (per_sequence_kl.sum() / full_attention_mask.shape[0]) + metrics.append((kl_div * attention_mask).sum() / full_attention_mask.sum().clamp(min=1.0)) + is_clipped = ((log_ratio < -epsilon_low) & (advantages.unsqueeze(1) < 0)) | ( + (log_ratio > epsilon_high) & (advantages.unsqueeze(1) > 0) + ) + metrics.append((is_clipped * attention_mask).sum() / full_attention_mask.sum().clamp(min=1.0)) + return loss, metrics + + class TorchLMHeadGRPO(torch.nn.Module): def __init__( self, @@ -392,6 +431,234 @@ def test_selective_chunk_forward_matches_reference(B, T, H, V, dtype, atol, rtol assert_verbose_allclose(out, ref, atol=atol, rtol=rtol) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 1e-5, 1e-5), + (torch.bfloat16, 2e-2, 5e-2), + ], +) +@pytest.mark.parametrize("chunk_size", [1, 2, 5]) +@pytest.mark.parametrize( + "beta, use_bias_correction_kl", + [(0.0, False), (0.04, False), (0.04, True)], +) +def test_gmpo_correctness_across_outer_chunks(dtype, atol, rtol, chunk_size, beta, use_bias_correction_kl): + set_seed() + B, T, H, V = 5, 7, 13, 31 + epsilon_low = 0.15 + epsilon_high = 0.4 + temperature = 0.9 + + inputs = torch.randn(B, T, H, device=device, dtype=dtype) + weight = torch.randn(V, H, device=device, dtype=dtype) * 0.2 + bias = torch.randn(V, device=device, dtype=dtype) * 0.1 + input1 = inputs.detach().clone().requires_grad_(True) + input2 = inputs.detach().clone().requires_grad_(True) + weight1 = weight.detach().clone().requires_grad_(True) + weight2 = weight.detach().clone().requires_grad_(True) + bias1 = bias.detach().clone().requires_grad_(True) + bias2 = bias.detach().clone().requires_grad_(True) + selected_token_ids = torch.randint(0, V, (B, T), device=device) + attention_mask = torch.tensor( + [ + [1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1, 0], + [1, 0, 0, 0, 0, 0, 0], + ], + device=device, + dtype=torch.float32, + ) + advantages = torch.tensor([-1.0, 0.0, 0.5, 1.25, -0.4], device=device) + + with torch.no_grad(): + initial_logits = (inputs @ weight.t() + bias) / temperature + initial_logps = ( + F.log_softmax(initial_logits.float(), dim=-1).gather(-1, selected_token_ids.unsqueeze(-1)).squeeze(-1) + ) + old_per_token_logps = initial_logps - torch.linspace(-0.8, 0.8, B * T, device=device).reshape(B, T) + ref_per_token_logps = initial_logps + 0.05 * torch.sin(torch.arange(B * T, device=device).reshape(B, T)) + + logits = (input1 @ weight1.t() + bias1) / temperature + per_token_logps = F.log_softmax(logits.float(), dim=-1).gather(-1, selected_token_ids.unsqueeze(-1)).squeeze(-1) + loss1, metrics1 = torch_gmpo_loss( + per_token_logps, + attention_mask, + advantages, + attention_mask, + old_per_token_logps, + ref_per_token_logps, + epsilon_low, + epsilon_high, + beta, + use_bias_correction_kl, + ) + loss2, metrics2 = LigerFusedLinearGRPOLoss( + beta=beta, + compiled=False, + use_ref_model=True, + chunk_size=chunk_size, + epsilon_low=epsilon_low, + epsilon_high=epsilon_high, + loss_type="gmpo", + temperature=temperature, + use_bias_correction_kl=use_bias_correction_kl, + )( + input2, + weight2, + selected_token_ids, + attention_mask, + advantages, + bias=bias2, + ref_per_token_logps=ref_per_token_logps, + old_per_token_logps=old_per_token_logps, + ) + + assert_verbose_allclose(loss1, loss2, atol=atol, rtol=rtol) + assert len(metrics1) == len(metrics2) == 1 + (1 if beta else 0) + for metric1, metric2 in zip(metrics1, metrics2): + assert_verbose_allclose(metric1, metric2, atol=atol, rtol=rtol) + + loss1.backward() + loss2.backward() + assert_verbose_allclose(input1.grad, input2.grad, atol=atol, rtol=rtol) + assert_verbose_allclose(weight1.grad, weight2.grad, atol=atol, rtol=rtol) + assert_verbose_allclose(bias1.grad, bias2.grad, atol=atol, rtol=rtol) + + +def test_gmpo_compiled_forward_backward(): + B, T, H, V = 2, 3, 5, 11 + inputs = torch.randn(B, T, H, device=device, requires_grad=True) + weight = torch.randn(V, H, device=device, requires_grad=True) + selected_token_ids = torch.randint(0, V, (B, T), device=device) + attention_mask = torch.tensor([[1, 1, 1], [1, 1, 0]], device=device) + advantages = torch.tensor([1.0, -1.0], device=device) + + loss, metrics = LigerFusedLinearGRPOLoss( + beta=0.0, + compiled=True, + use_ref_model=False, + loss_type="gmpo", + )(inputs, weight, selected_token_ids, attention_mask, advantages) + loss.backward() + + assert torch.isfinite(loss) + assert len(metrics) == 1 + assert torch.isfinite(metrics[0]) + assert torch.isfinite(inputs.grad).all() + assert torch.isfinite(weight.grad).all() + + +def test_gmpo_log_clip_boundaries_match_reference_forward_backward(): + dtype = torch.float64 + lower_bound = torch.tensor(-0.15, device=device, dtype=dtype) + upper_bound = torch.tensor(0.4, device=device, dtype=dtype) + negative_infinity = torch.tensor(float("-inf"), device=device, dtype=dtype) + positive_infinity = torch.tensor(float("inf"), device=device, dtype=dtype) + log_ratio = torch.stack( + ( + torch.stack( + ( + upper_bound, + torch.nextafter(upper_bound, positive_infinity), + torch.nextafter(upper_bound, negative_infinity), + ) + ), + torch.stack( + ( + lower_bound, + torch.nextafter(lower_bound, negative_infinity), + torch.nextafter(lower_bound, positive_infinity), + ) + ), + ) + ) + per_token_logps_ref = torch.zeros_like(log_ratio, requires_grad=True) + per_token_logps = per_token_logps_ref.detach().clone().requires_grad_(True) + old_per_token_logps = -log_ratio + attention_mask = torch.ones_like(log_ratio) + advantages = torch.tensor([1.0, -1.0], device=device, dtype=dtype) + + expected_loss, expected_metrics = torch_gmpo_loss( + per_token_logps_ref, + attention_mask, + advantages, + attention_mask, + old_per_token_logps, + per_token_logps_ref.detach(), + epsilon_low=0.15, + epsilon_high=0.4, + beta=0.0, + ) + loss, metrics = LigerFusedLinearGRPOFunction.ppo_loss_fn( + per_token_logps, + attention_mask, + advantages, + attention_mask, + old_per_token_logps=old_per_token_logps, + epsilon_low=0.15, + epsilon_high=0.4, + beta=0.0, + loss_type="gmpo", + ) + expected_loss.backward() + loss.backward() + + torch.testing.assert_close(loss, expected_loss, rtol=0.0, atol=0.0) + torch.testing.assert_close(per_token_logps.grad, per_token_logps_ref.grad, rtol=0.0, atol=0.0) + assert len(metrics) == len(expected_metrics) == 1 + torch.testing.assert_close(metrics[0], expected_metrics[0], rtol=0.0, atol=0.0) + torch.testing.assert_close(metrics[0], torch.tensor(1 / 3, device=device, dtype=dtype), rtol=0.0, atol=0.0) + + +def test_gmpo_extreme_active_clips_are_finite(): + per_token_logps_ref = torch.zeros(3, 3, device=device, requires_grad=True) + per_token_logps = per_token_logps_ref.detach().clone().requires_grad_(True) + log_ratio = torch.tensor( + [[1000.0, 100.0, 0.1], [-1000.0, -100.0, -0.1], [1000.0, 100.0, 0.1]], + device=device, + ) + old_per_token_logps = per_token_logps.detach() - log_ratio + attention_mask = torch.ones_like(per_token_logps) + advantages = torch.tensor([1.0, -1.0, 0.0], device=device) + + expected_loss, expected_metrics = torch_gmpo_loss( + per_token_logps_ref, + attention_mask, + advantages, + attention_mask, + old_per_token_logps, + per_token_logps_ref.detach(), + epsilon_low=0.15, + epsilon_high=0.4, + beta=0.0, + ) + + loss, metrics = LigerFusedLinearGRPOFunction.ppo_loss_fn( + per_token_logps, + attention_mask, + advantages, + attention_mask, + old_per_token_logps=old_per_token_logps, + epsilon_low=0.15, + epsilon_high=0.4, + beta=0.0, + loss_type="gmpo", + ) + expected_loss.backward() + loss.backward() + + assert torch.isfinite(loss) + assert torch.isfinite(per_token_logps.grad).all() + assert_verbose_allclose(loss, expected_loss) + assert_verbose_allclose(per_token_logps.grad, per_token_logps_ref.grad) + assert len(metrics) == len(expected_metrics) == 1 + assert_verbose_allclose(metrics[0], expected_metrics[0]) + assert_verbose_allclose(metrics[0], torch.tensor(4 / 9, device=device)) + + @pytest.mark.parametrize("loss_type", ["dapo", "grpo"]) @pytest.mark.parametrize("compiled", [True, False]) def test_correctness_large_seq_exercises_chunking(loss_type, compiled): @@ -1086,6 +1353,86 @@ def test_reduce_grpo_loss_requires_max_completion_length(): assert_verbose_allclose(reduced, expected) +def test_gmpo_rejects_sequence_importance_sampling(): + B, T, H, V = 2, 4, 8, 16 + loss_fn = LigerFusedLinearGRPOLoss( + beta=0.0, + compiled=False, + use_ref_model=False, + loss_type="gmpo", + importance_sampling_level="sequence", + ) + inputs = torch.randn(B, T, H, device=device, requires_grad=True) + weight = torch.randn(V, H, device=device, requires_grad=True) + 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) + + with pytest.raises( + ValueError, + match="Sequence-level importance sampling is not supported for loss_type='gmpo'", + ): + loss_fn( + inputs, + weight, + selected_token_ids, + attention_mask, + advantages, + ) + + +@pytest.mark.parametrize("option", ["vllm", "delta"]) +def test_gmpo_ignores_vllm_ratio_and_delta(option): + set_seed() + B, T, H, V = 2, 4, 8, 16 + inputs = torch.randn(B, T, H, device=device) + weight = torch.randn(V, H, device=device) + bias = torch.randn(V, device=device) + selected_token_ids = torch.randint(0, V, (B, T), device=device) + attention_mask = torch.tensor([[1, 1, 1, 1], [1, 1, 0, 0]], device=device) + advantages = torch.tensor([1.0, -1.0], device=device) + with torch.no_grad(): + current_logps = LigerFusedLinearPPOBase.chunk_forward(inputs, weight, selected_token_ids, bias=bias) + old_per_token_logps = current_logps - torch.linspace(-0.6, 0.6, B * T, device=device).reshape(B, T) + + def run(with_option): + module_kwargs = {"delta": 0.5} if option == "delta" and with_option else {} + forward_kwargs = ( + {"vllm_is_ratio": torch.full((B, T), 0.5, device=device)} if option == "vllm" and with_option else {} + ) + input_tensor = inputs.detach().clone().requires_grad_(True) + weight_tensor = weight.detach().clone().requires_grad_(True) + bias_tensor = bias.detach().clone().requires_grad_(True) + loss, metrics = LigerFusedLinearGRPOLoss( + beta=0.0, + compiled=False, + use_ref_model=False, + loss_type="gmpo", + **module_kwargs, + )( + input_tensor, + weight_tensor, + selected_token_ids, + attention_mask, + advantages, + bias=bias_tensor, + old_per_token_logps=old_per_token_logps, + **forward_kwargs, + ) + loss.backward() + return loss.detach(), metrics, input_tensor.grad, weight_tensor.grad, bias_tensor.grad + + baseline = run(False) + with_option = run(True) + for expected, actual in zip(baseline, with_option): + if isinstance(expected, tuple): + assert len(expected) == len(actual) + for expected_metric, actual_metric in zip(expected, actual): + assert_verbose_allclose(expected_metric, actual_metric) + else: + assert_verbose_allclose(expected, actual) + + @pytest.mark.parametrize("loss_type", ["cispo", "sapo"]) def test_sequence_level_rejects_unsupported_loss_types(loss_type): """Sequence-level importance sampling should raise ValueError for cispo and sapo.""" From c7b8c79dcdcf63b19ab100d717e8a4815bdfc918 Mon Sep 17 00:00:00 2001 From: Chino Date: Wed, 29 Jul 2026 19:45:10 +0800 Subject: [PATCH 2/3] Update src/liger_kernel/chunked_loss/grpo_loss.py Co-authored-by: Tcc0403 <76503978+Tcc0403@users.noreply.github.com> --- src/liger_kernel/chunked_loss/grpo_loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/liger_kernel/chunked_loss/grpo_loss.py b/src/liger_kernel/chunked_loss/grpo_loss.py index 087f981a5..b9aee6e9d 100644 --- a/src/liger_kernel/chunked_loss/grpo_loss.py +++ b/src/liger_kernel/chunked_loss/grpo_loss.py @@ -79,8 +79,8 @@ def clip_coef_fn(coef, epsilon_low, epsilon_high, loss_type, advantages=None): positive_advantages = advantages.unsqueeze(1) >= 0 clipped_coef = torch.where( positive_advantages, - torch.minimum(coef, torch.as_tensor(epsilon_high, device=coef.device, dtype=coef.dtype)), - torch.maximum(coef, torch.as_tensor(-epsilon_low, device=coef.device, dtype=coef.dtype)), + torch.clamp(coef, max=epsilon_high), + torch.clamp(coef, min=-epsilon_low), ) is_lower_clipped = coef < -epsilon_low is_upper_clipped = coef > epsilon_high From 1d9971e4246030635c621f16cd4878b14c8f91b0 Mon Sep 17 00:00:00 2001 From: chinoll Date: Wed, 29 Jul 2026 19:55:31 +0800 Subject: [PATCH 3/3] Restore GRPO loss comments --- src/liger_kernel/chunked_loss/grpo_loss.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/liger_kernel/chunked_loss/grpo_loss.py b/src/liger_kernel/chunked_loss/grpo_loss.py index b9aee6e9d..ab3321ec5 100644 --- a/src/liger_kernel/chunked_loss/grpo_loss.py +++ b/src/liger_kernel/chunked_loss/grpo_loss.py @@ -159,6 +159,8 @@ def ppo_loss_fn( "and 'sequence'." ) + # From here, log_importance_weights (and all subsequent tensors, coef_1, coef_2, etc.) shape depends on + # importance_sampling_level: "token" level: (B, T); "sequence" level: (B, 1) # GMPO coefficients are log-domain; all other loss types use ratio-domain coefficients. coef_1 = log_importance_weights if loss_type == "gmpo" else torch.exp(log_importance_weights) coef_2, is_lower_clipped, is_upper_clipped = clip_coef_fn( @@ -230,6 +232,7 @@ def ppo_loss_fn( # importance_sampling_level (token: (B, T); sequence: (B, 1)), # mirroring TRL's ``per_token_kl * coef_1`` (un-clamped, before delta). kl_div = kl_div * torch.exp(log_importance_weights) + # Combine losses per_token_loss = per_token_loss + beta * kl_div # Note: We normalize by the number of tokens in the batch (using full_attention_mask),