perf(grpo): skip the redundant pre-update policy forward when it cannot differ

swift/rlhf_trainers/grpo_trainer.py

@@ -758,9 +758,10 @@ def _prepare_batch_inputs(self, samples: List[GRPOSample]) -> List[Dict[str, Any
             origin_data = batch if (self.dynamic_num_samples and self.is_multimodal) else None
 
             with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs):
-                grpo_batch.old_per_token_logps = (
-                    self._get_per_token_logps_and_entropies(
-                        self.model, model_inputs, grpo_batch, origin_data=origin_data)[0])
+                if self._need_old_per_token_logps(grpo_batch):
+                    grpo_batch.old_per_token_logps = (
+                        self._get_per_token_logps_and_entropies(
+                            self.model, model_inputs, grpo_batch, origin_data=origin_data)[0])
                 if self.beta == 0.0:
                     ref_per_token_logps = None
                 elif self.ref_model is not None:
@@ -1831,6 +1832,20 @@ def training_step(self, model: nn.Module, inputs: DataType, num_items_in_batch=N
             self.eval_flag = False
         return super().training_step(model, inputs, num_items_in_batch)
 
+    def _need_old_per_token_logps(self, grpo_batch) -> bool:
+        """Whether the pre-update policy logps need a forward of their own.
+
+        With the default setup (`num_iterations=1` and the rollout batch consumed inside a single
+        optimizer step) the pre-update policy and the current policy are the same weights, so the
+        logps produced by this extra forward are bit-for-bit identical to
+        `per_token_logps.detach()` — the forward is pure overhead. It is only worth computing when
+        the weights can actually have moved, or when something reads the field directly without the
+        `per_token_logps.detach()` fallback used at the loss site.
+        """
+        return bool(self.old_policy() or (self.kl_in_reward and self.beta != 0.0)
+                    or grpo_batch.teacher_per_token_logps is not None or grpo_batch.rollout_per_token_logps is not None
+                    or self.args.async_generate or self.args.use_liger_kernel)
+
     def old_policy(self):
         if self.template.sequence_parallel_size == 1:
             return (self.num_iterations > 1

tests/general/test_grpo_old_logps_gate.py

@@ -0,0 +1,61 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import unittest
+from types import SimpleNamespace
+
+from swift.rlhf_trainers.grpo_trainer import GRPOTrainer
+
+_need_old_per_token_logps = GRPOTrainer._need_old_per_token_logps
+
+
+def make_trainer(**kwargs):
+    trainer = dict(
+        old_policy=lambda: False,
+        kl_in_reward=False,
+        beta=0.0,
+        args=SimpleNamespace(async_generate=False, use_liger_kernel=False),
+    )
+    trainer.update(kwargs)
+    return SimpleNamespace(**trainer)
+
+
+def make_batch(**kwargs):
+    batch = dict(teacher_per_token_logps=None, rollout_per_token_logps=None)
+    batch.update(kwargs)
+    return SimpleNamespace(**batch)
+
+
+class TestNeedOldPerTokenLogps(unittest.TestCase):
+    """The extra forward over the pre-update policy is only skipped when it is provably redundant."""
+
+    def test_default_setup_skips_the_forward(self):
+        # num_iterations=1, single optimizer step, no reader without a fallback.
+        self.assertFalse(_need_old_per_token_logps(make_trainer(), make_batch()))
+
+    def test_multiple_iterations_keep_the_forward(self):
+        self.assertTrue(_need_old_per_token_logps(make_trainer(old_policy=lambda: True), make_batch()))
+
+    def test_kl_in_reward_keeps_the_forward_only_with_beta(self):
+        # kl_in_reward computes its KL straight from the field, with no `per_token_logps` fallback.
+        self.assertTrue(_need_old_per_token_logps(make_trainer(kl_in_reward=True, beta=0.1), make_batch()))
+        self.assertFalse(_need_old_per_token_logps(make_trainer(kl_in_reward=True, beta=0.0), make_batch()))
+
+    def test_teacher_kl_keeps_the_forward(self):
+        batch = make_batch(teacher_per_token_logps=object())
+        self.assertTrue(_need_old_per_token_logps(make_trainer(), batch))
+
+    def test_rollout_logps_keep_the_forward(self):
+        # The Liger path reads the field directly to build the rollout importance ratio.
+        batch = make_batch(rollout_per_token_logps=object())
+        self.assertTrue(_need_old_per_token_logps(make_trainer(), batch))
+
+    def test_async_generate_keeps_the_forward(self):
+        trainer = make_trainer(args=SimpleNamespace(async_generate=True))
+        self.assertTrue(_need_old_per_token_logps(trainer, make_batch()))
+
+    def test_liger_kernel_keeps_the_forward(self):
+        trainer = make_trainer(args=SimpleNamespace(async_generate=False, use_liger_kernel=True))
+        self.assertTrue(_need_old_per_token_logps(trainer, make_batch()))
+
+
+if __name__ == '__main__':
+    unittest.main()