From 1912a94eb814bce94bd509225923c10ba7339a92 Mon Sep 17 00:00:00 2001
From: taking-lying-flat <1615405@qq.com>
Date: Sun, 20 Sep 2026 20:14:44 +0800
Subject: [PATCH] fix: weight PPO token-mean micro-batches by valid tokens

---
 src/twinkle/loss/grpo.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py
index e146edee..46161468 100644
--- a/src/twinkle/loss/grpo.py
+++ b/src/twinkle/loss/grpo.py
@@ -384,6 +384,16 @@ def __init__(self, loss_agg_mode: Literal['token-mean', 'seq-mean-token-mean'] =
                              f'Expected one of {sorted(self._LOSS_AGG_MODES)}.')
         self.loss_agg_mode = loss_agg_mode
 
+    def micro_batch_scale(self, inputs, indices):
+        if self.loss_agg_mode != 'token-mean':
+            return super().micro_batch_scale(inputs, indices)
+        import torch
+        token_counts = [
+            self._resolve_loss_mask(model_input, torch.atleast_2d(torch.as_tensor(model_input['labels']))).sum().item()
+            for model_input in inputs
+        ]
+        return sum(token_counts[index] for index in indices) / max(sum(token_counts), 1)
+
     def _aggregate_loss(
         self,
         per_token_loss: 'torch.Tensor',