fix: weight PPO token-mean micro-batches by valid tokens

[taking-lying-flat]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

PPOLoss(loss_agg_mode='token-mean') averages valid tokens within a micro-batch but inherits sample-count weighting from GRPO. With response lengths 1 and 3 and advantages +1 and -1, changing micro_batch_size from 2 to 1 changes the loss from 0.5 to 0 and the same binary policy parameter's gradient from 0.25 to 0.

Override PPO's micro-batch weight with the fraction of valid tokens in the logical input batch. Reuse _resolve_loss_mask() so ignored labels and completion_mask match the loss's scoring mask. Keep sample-count weighting for seq-mean-token-mean and preserve the existing loss-return and accumulation contracts.

Experiment results

• Reproduced the failure with float64 using the real Transformers micro-batch planner, loss calculation, backward and gradient normalization. After the fix, both micro-batch sizes give loss 0.5 and gradient 0.25.
• 316 local integration/weight checks passed, covering static/FFD/KK batching, padded/packed inputs, completion masks, zero valid tokens, both aggregation modes, gradient accumulation, float32/float64 and list/NumPy/tensor inputs.
• 45 existing tests passed: tests/loss/test_ppo.py, tests/loss/test_grpo_gkd.py, tests/loss/test_bnpo_token_mean.py and tests/model/test_micro_batch.py.
• pre-commit run --all-files and git diff --check passed.

Only the production change in src/twinkle/loss/grpo.py is included (+10 lines).