[GRPO] Multiply the KL term by the per-token importance sampling ratio

tests/test_grpo_trainer.py

@@ -1952,6 +1952,60 @@ def test_train_bias_correction_kl(self, use_liger_kernel, use_bias_correction_kl
             new_param = trainer.model.get_parameter(n)
             assert not torch.equal(param, new_param), f"Parameter {n} has not changed."
 
+    def test_bias_correction_kl_uses_the_per_token_ratio_at_sequence_level(self):
+        # Regression test for #6586. The KL term is per-token, so its importance sampling correction must use the
+        # per-token ratio even when importance_sampling_level="sequence"; broadcasting the sequence-level weight onto
+        # the per-token KL gives a different gradient. With zero advantages the loss is the KL term alone, so it must
+        # not depend on the importance sampling level.
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+        training_args = GRPOConfig(
+            output_dir=self.tmp_dir,
+            loss_type="grpo",
+            importance_sampling_level="sequence",
+            beta=0.1,
+            use_bias_correction_kl=True,
+            report_to="none",
+        )
+        trainer = GRPOTrainer(
+            model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
+            reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+            args=training_args,
+            train_dataset=dataset,
+        )
+        trainer.model.eval()
+
+        device = next(trainer.model.parameters()).device
+        batch_size, prompt_len, completion_len = 2, 3, 4
+        prompt_ids = torch.randint(1, 1000, (batch_size, prompt_len), device=device)
+        completion_ids = torch.randint(1, 1000, (batch_size, completion_len), device=device)
+        prompt_mask = torch.ones_like(prompt_ids)
+        completion_mask = torch.ones_like(completion_ids)
+        with torch.no_grad():
+            per_token_logps, _, _ = trainer._get_per_token_logps_and_entropies(
+                trainer.model,
+                torch.cat([prompt_ids, completion_ids], dim=1),
+                torch.cat([prompt_mask, completion_mask], dim=1),
+                completion_len,
+            )
+        # Per-token log-ratios that cancel out over each sequence: the sequence-level ratio is exactly 1, so a
+        # sequence-level correction leaves the KL term unchanged while the per-token correction does not.
+        log_ratio = torch.tensor([[0.5, -0.5, 0.5, -0.5]] * batch_size, device=device)
+        inputs = {
+            "prompt_ids": prompt_ids,
+            "prompt_mask": prompt_mask,
+            "completion_ids": completion_ids,
+            "completion_mask": completion_mask,
+            "advantages": torch.zeros(batch_size, device=device),
+            "old_per_token_logps": per_token_logps - log_ratio,
+            "ref_per_token_logps": per_token_logps + 1.0,
+        }
+
+        sequence_level_loss = trainer._compute_loss(trainer.model, inputs)
+        trainer.importance_sampling_level = "token"
+        token_level_loss = trainer._compute_loss(trainer.model, inputs)
+
+        torch.testing.assert_close(sequence_level_loss, token_level_loss)
+
     def test_reward_func_wrong_number_of_rewards(self):
         # A reward function that returns the wrong number of rewards should raise a clear error instead of silently
         # broadcasting (when it returns a single value) or failing later with an opaque shape error.

trl/trainer/grpo_config.py

@@ -348,12 +348,11 @@ class GRPOConfig(_BaseConfig):
             parameter corresponds to the `delta` threshold in Equation 9 of the [DeepSeek-V3.2
             paper](https://huggingface.co/papers/2512.02556). It expects a positive value (e.g., 0.5).
         use_bias_correction_kl (`bool`, *optional*, defaults to `True`):
-            Whether to multiply the KL term by the importance sampling ratio, so that the KL gradient becomes the
-            unbiased reverse-KL gradient, as described in the [DeepSeek-V3.2
-            paper](https://huggingface.co/papers/2512.02556). This changes the KL gradient whenever `beta != 0`,
-            including on-policy: the ratio is differentiable, so it affects the gradient even where its value is
-            exactly 1. The unbiased reverse-KL property holds for `importance_sampling_level="token"`; with
-            `"sequence"` a sequence-level weight is broadcast onto the per-token KL.
+            Whether to multiply the KL term by the per-token importance sampling ratio, so that the KL gradient becomes
+            the unbiased reverse-KL gradient, as described in the [DeepSeek-V3.2
+            paper](https://huggingface.co/papers/2512.02556). The per-token ratio is used regardless of
+            `importance_sampling_level`. This changes the KL gradient whenever `beta != 0`, including on-policy: the
+            ratio is differentiable, so it affects the gradient even where its value is exactly 1.
 
         > Parameters that control the logging
 
@@ -967,12 +966,11 @@ class GRPOConfig(_BaseConfig):
     use_bias_correction_kl: bool = field(
         default=True,
         metadata={
-            "help": "Whether to multiply the KL term by the importance sampling ratio, so that the KL gradient "
-            "becomes the unbiased reverse-KL gradient, as described in the [DeepSeek-V3.2 "
-            "paper](https://huggingface.co/papers/2512.02556). This changes the KL gradient whenever `beta != 0`, "
-            "including on-policy: the ratio is differentiable, so it affects the gradient even where its value is "
-            "exactly 1. The unbiased reverse-KL property holds for `importance_sampling_level='token'`; with "
-            "'sequence' a sequence-level weight is broadcast onto the per-token KL."
+            "help": "Whether to multiply the KL term by the per-token importance sampling ratio, so that the KL "
+            "gradient becomes the unbiased reverse-KL gradient, as described in the [DeepSeek-V3.2 "
+            "paper](https://huggingface.co/papers/2512.02556). The per-token ratio is used regardless of "
+            "`importance_sampling_level`. This changes the KL gradient whenever `beta != 0`, including on-policy: "
+            "the ratio is differentiable, so it affects the gradient even where its value is exactly 1."
         },
     )
 

trl/trainer/grpo_trainer.py

@@ -3193,9 +3193,10 @@ def _compute_loss(self, model, inputs):
             per_token_kl = (
                 torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
             )
-            # Importance sampling correction for the KL divergence
+            # Importance sampling correction for the KL divergence. The KL is per-token, so the correction uses the
+            # per-token ratio π_θ/π_old regardless of `importance_sampling_level`.
             if self.args.use_bias_correction_kl:
-                per_token_kl = per_token_kl * coef_1
+                per_token_kl = per_token_kl * torch.exp(log_ratio)
 
         # 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)