[bugfix] fix grpo 'num_iterms_in_batch' in forward kwargs (#6717)

hjh0119 · web-flow · commit 032993c44b10 · 2025-11-23T19:26:37.000+08:00
diff --git a/swift/trainers/rlhf_trainer/grpo_trainer.py b/swift/trainers/rlhf_trainer/grpo_trainer.py
@@ -1314,7 +1314,7 @@ def _get_per_token_logps_and_entropies_sp(
             k: v
             for k, v in inputs.items() if k not in [
                 'logits_to_keep', 'completion_mask', 'ref_per_token_logps', 'advantages', 'old_per_token_logps',
-                'truncated_mask', 'seq_lengths'
+                'truncated_mask', 'seq_lengths', 'num_items_in_batch'
             ]
         }
         sequence_parallel.prepare_inputs(inputs)
@@ -1394,7 +1394,7 @@ def _get_per_token_logps_and_entropies_single(self,
                 k: v
                 for k, v in inputs.items() if k not in [
                     'logits_to_keep', 'completion_mask', 'ref_per_token_logps', 'advantages', 'old_per_token_logps',
-                    'truncated_mask', 'seq_lengths'
+                    'truncated_mask', 'seq_lengths', 'num_items_in_batch'
                 ]
             }
             if 'logits_to_keep' in self.model_kwarg_keys:
@@ -1488,7 +1488,7 @@ def _get_last_hidden_state(self, unwrapped_model, inputs, logits_to_keep):
                 k: v
                 for k, v in inputs.items() if k not in [
                     'logits_to_keep', 'completion_mask', 'ref_per_token_logps', 'advantages', 'old_per_token_logps',
-                    'truncated_mask', 'seq_lengths'
+                    'truncated_mask', 'seq_lengths', 'num_items_in_batch'
                 ]
             }
             if 'logits_to_keep' in self.model_kwarg_keys:

Original file line number	Diff line number	Diff line change
`@@ -1314,7 +1314,7 @@ def _get_per_token_logps_and_entropies_sp(`
`1314`	`1314`	`k: v`
`1315`	`1315`	`for k, v in inputs.items() if k not in [`
`1316`	`1316`	`'logits_to_keep', 'completion_mask', 'ref_per_token_logps', 'advantages', 'old_per_token_logps',`
`1317`		`- 'truncated_mask', 'seq_lengths'`
	`1317`	`+ 'truncated_mask', 'seq_lengths', 'num_items_in_batch'`
`1318`	`1318`	`]`
`1319`	`1319`	`}`
`1320`	`1320`	`sequence_parallel.prepare_inputs(inputs)`
`@@ -1394,7 +1394,7 @@ def _get_per_token_logps_and_entropies_single(self,`
`1394`	`1394`	`k: v`
`1395`	`1395`	`for k, v in inputs.items() if k not in [`
`1396`	`1396`	`'logits_to_keep', 'completion_mask', 'ref_per_token_logps', 'advantages', 'old_per_token_logps',`
`1397`		`- 'truncated_mask', 'seq_lengths'`
	`1397`	`+ 'truncated_mask', 'seq_lengths', 'num_items_in_batch'`
`1398`	`1398`	`]`
`1399`	`1399`	`}`
`1400`	`1400`	`if 'logits_to_keep' in self.model_kwarg_keys:`
`@@ -1488,7 +1488,7 @@ def _get_last_hidden_state(self, unwrapped_model, inputs, logits_to_keep):`
`1488`	`1488`	`k: v`
`1489`	`1489`	`for k, v in inputs.items() if k not in [`
`1490`	`1490`	`'logits_to_keep', 'completion_mask', 'ref_per_token_logps', 'advantages', 'old_per_token_logps',`
`1491`		`- 'truncated_mask', 'seq_lengths'`
	`1491`	`+ 'truncated_mask', 'seq_lengths', 'num_items_in_batch'`
`1492`	`1492`	`]`
`1493`	`1493`	`}`
`1494`	`1494`	`if 'logits_to_keep' in self.model_kwarg_keys:`