fix loss micro_batch_scale

src/twinkle/loss/grpo.py

@@ -4,6 +4,7 @@
 
 from twinkle.data_format import LossOutput
 from twinkle.loss.base import Loss
+from twinkle.loss.utils import count_loss_tokens
 from twinkle.utils.torch_utils import selective_log_softmax
 
 if TYPE_CHECKING:
@@ -632,6 +633,7 @@ def _configured_grpo(
                 channel_loss['entropy'] = torch.stack(
                     ((entropies * alignment_mask).sum().detach().float(), count))
             if self.top_entropy_quantile < 1.0 and alignment_mask.any():
+                # The entropy threshold is local to the current micro-batch.
                 threshold = torch.quantile(
                     entropies[alignment_mask].float(), 1.0 - self.top_entropy_quantile)
                 entropy_mask = entropies.ge(threshold)
@@ -793,6 +795,17 @@ 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)
+        # The configured path scores labels alone; the base path also intersects completion_mask.
+        loss_mask_fn = None if self._advanced else self._resolve_loss_mask
+        token_counts = [
+            count_loss_tokens(row, self.ignore_index, loss_mask_fn=loss_mask_fn) for row in inputs
+        ]
+        total_tokens = sum(token_counts)
+        return sum(token_counts[index] for index in indices) / total_tokens if total_tokens else 0.0
+
     def _aggregate_loss(
         self,
         per_token_loss: 'torch.Tensor',

src/twinkle/loss/infonce.py

@@ -169,6 +169,8 @@ class InfonceLoss(EmbeddingLoss):
     such group. Setting ``use_batch=True`` enables in-batch negatives and,
     when distributed is initialized, gathers embeddings from all DP ranks
     (only the local shard keeps gradients).
+    The negative pool covers the current micro-batch and its DP peers, not other
+    micro-batches in the gradient-accumulation window.
 
     Args:
         temperature: Logit scaling factor.
@@ -454,6 +456,7 @@ class OnlineContrastiveLoss(EmbeddingLoss):
     dropping them concentrates the gradient on the decision boundary. The single-element
     fallbacks (``negs.mean()`` / ``poss.mean()``) mirror legacy: with fewer than two pairs on a
     side there is no counterpart to compare against.
+    Hard-pair selection is performed independently within each micro-batch.
 
     Args:
         margin: Hinge for the selected negative pairs.

src/twinkle/loss/utils.py

@@ -1,12 +1,19 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 
 
-def count_loss_tokens(model_input, ignore_index: int) -> int:
-    """Count non-ignored labels in one sequence or a nested packing unit."""
+def count_loss_tokens(model_input, ignore_index: int, *, loss_mask_fn=None) -> int:
+    """Count scored tokens in one sequence or packing unit, optionally using the loss's own mask."""
     # A packing unit stays intact until the processor expands it into sequences.
     if isinstance(model_input, (list, tuple)):
-        return sum(count_loss_tokens(item, ignore_index) for item in model_input)
+        return sum(count_loss_tokens(item, ignore_index, loss_mask_fn=loss_mask_fn) for item in model_input)
     labels = model_input["labels"]
+    if loss_mask_fn is not None:
+        import torch
+
+        labels = torch.as_tensor(labels)
+        if labels.dim() == 1:
+            labels = labels.unsqueeze(0)
+        return int(loss_mask_fn(model_input, labels).sum().item())
     if hasattr(labels, "ne"):
         return int(labels.ne(ignore_index).sum().item())
     return sum(int(token != ignore_index) for token in labels)

src/twinkle/loss/value.py

@@ -5,6 +5,7 @@
 from twinkle.data_format import LossOutput
 from twinkle.loss.base import Loss
 from twinkle.loss.grpo import GRPOLoss
+from twinkle.loss.utils import count_loss_tokens
 
 if TYPE_CHECKING:
     import torch
@@ -28,7 +29,9 @@ def __init__(self, cliprange_value: float = 0.2, vf_coef: float = 1.0, ignore_in
         self._aligner = GRPOLoss(ignore_index=ignore_index)
 
     def micro_batch_scale(self, inputs, indices):
-        return len(indices) / len(inputs)
+        token_counts = [count_loss_tokens(row, self.ignore_index) for row in inputs]
+        total_tokens = sum(token_counts)
+        return sum(token_counts[index] for index in indices) / total_tokens if total_tokens else 0.0
 
     def __call__(
         self,

src/twinkle/model/megatron/megatron.py

@@ -509,7 +509,6 @@ def forward_backward(self,
             variable_seq_lengths=self.variable_seq_lengths,
             attention_mask_type=attention_mask_type,
         )
-
         # Get parallelism settings for sequence padding and splitting
         cp_size = self.device_mesh.cp_world_size
         # MoE routing replay orchestration of this call (a null object when disabled or mode-less);

tests/loss/test_ppo.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 from twinkle.advantage import GAEAdvantage
@@ -29,6 +30,30 @@ def test_ppo_loss_aggregation_modes():
     torch.testing.assert_close(sequence_mean, torch.tensor(3.0))
 
 
+def test_ppo_micro_batch_scale_follows_each_loss_denominator():
+    rows = [
+        {'labels': [-100, 1, 2, 3], 'completion_mask': [0, 1, 0, 0]},
+        {'labels': [-100, 1, 2, 3], 'completion_mask': [0, 1, 1, 1]},
+    ]
+    assert PPOLoss().micro_batch_scale(rows, [0]) == pytest.approx(1 / 4)
+    assert PPOLoss(loss_agg_mode='seq-mean-token-mean').micro_batch_scale(rows, [0]) == pytest.approx(1 / 2)
+    # The critic scores all trainable labels, independently of the policy's completion_mask.
+    assert PPOValueLoss().micro_batch_scale(rows, [0]) == pytest.approx(1 / 2)
+    # Advanced GRPO/PPO controls use labels alone, matching _configured_grpo.
+    assert PPOLoss(log_entropy=True).micro_batch_scale(rows, [0]) == pytest.approx(1 / 2)
+
+
+@pytest.mark.parametrize('factory', [PPOLoss, PPOValueLoss])
+def test_ppo_micro_batch_scale_counts_packing_units_and_custom_ignore_index(factory):
+    rows = [[{'labels': [0, 1]}, {'labels': [2, 3]}], [{'labels': [0, 4]}]]
+    assert factory(ignore_index=0).micro_batch_scale(rows, [0]) == pytest.approx(3 / 4)
+
+
+@pytest.mark.parametrize('factory', [PPOLoss, PPOValueLoss])
+def test_ppo_micro_batch_scale_handles_no_scored_tokens(factory):
+    assert factory().micro_batch_scale([{'labels': [-100]}, {'labels': [-100]}], [0]) == 0.
+
+
 def test_value_loss_without_clipping():
     values = torch.tensor([[0.0, 1.0, 10.0]], requires_grad=True)
     result = PPOValueLoss(epsilon=0.2)(

tests/model/test_batch_context_loss.py

@@ -0,0 +1,158 @@
+"""Batch-dependent objectives use the statistics of each physical micro-batch."""
+import sys
+from inspect import unwrap
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from twinkle.loss import GRPOLoss
+from twinkle.loss.clip_cov import ClipCovLoss
+from twinkle.loss.grpo import BNPOLoss, CISPOLoss, PPOLoss
+from twinkle.loss.infonce import EmbeddingLoss, InfonceLoss, OnlineContrastiveLoss
+from twinkle.loss.kl_cov import KLCovLoss
+from twinkle.model.optimizer_group import BaseOptimizerGroup
+from twinkle.processor import InputProcessor
+from .test_grpo_micro_batch import TrainingHarness, data
+
+BATCH_CONTEXT_LOSS_FACTORIES = [
+    pytest.param(InfonceLoss, id='infonce-batch-negatives'),
+    pytest.param(OnlineContrastiveLoss, id='online-contrastive'),
+    pytest.param(lambda: ClipCovLoss(clip_cov_ratio=.5, clip_cov_lb=.001, clip_cov_ub=.5), id='clip-cov'),
+    pytest.param(lambda: KLCovLoss(kl_cov_ratio=.5), id='kl-cov'),
+    pytest.param(lambda: GRPOLoss(top_entropy_quantile=.5), id='grpo-entropy-filter'),
+    pytest.param(lambda: PPOLoss(top_entropy_quantile=.5), id='ppo-entropy-filter'),
+    pytest.param(lambda: CISPOLoss(top_entropy_quantile=.5), id='cispo-entropy-filter'),
+    pytest.param(lambda: BNPOLoss(top_entropy_quantile=.5, token_mean_scope='global'), id='bnpo-entropy-filter'),
+]
+
+
+class BatchContextHarness(TrainingHarness):
+    def forward(self, *, inputs, **kwargs):
+        super().forward(inputs=inputs, **kwargs)
+        group = self.optimizer_group['default']
+        status = group.train_status
+        status.outputs['entropies'] = 1. + status.outputs['logps'].square()
+        if isinstance(group.loss_instance, EmbeddingLoss):
+            embeddings = []
+            infonce = isinstance(group.loss_instance, InfonceLoss)
+            for row in inputs:
+                index = row['input_ids'][0] + 1
+                # Each input row is a complete query/candidate group or labelled pair.
+                vectors = torch.tensor([[1., .1 * index, .2], [.1 * index, 1., -.3], [-.4, .3, .1 * index]],
+                                       dtype=self.weight.dtype)
+                if not infonce:
+                    vectors = vectors[:2]
+                embeddings.append(F.normalize(vectors + self.weight * vectors.roll(1, dims=-1), dim=-1))
+            status.outputs['embeddings'] = torch.cat(embeddings)
+            status.inputs['labels'] = torch.tensor(
+                [1, 0, 0] * len(inputs) if infonce else [row['input_ids'][0] % 2 for row in inputs])
+        return dict(status.outputs)
+
+
+@pytest.mark.parametrize('factory', BATCH_CONTEXT_LOSS_FACTORIES)
+@pytest.mark.parametrize('options', [
+    pytest.param({}, id='unsplit'),
+    pytest.param({'micro_batch_size': 4}, id='fixed'),
+    pytest.param({'micro_batch_size': 6}, id='fixed-short-tail'),
+    pytest.param({'micro_batch_size': 4, 'dynamic_batching': True,
+                  'max_tokens_per_micro_batch': 18, 'packing_algorithm': 'ffd'}, id='dynamic-ffd'),
+    pytest.param({'micro_batch_size': 4, 'dynamic_batching': True,
+                  'max_tokens_per_micro_batch': 18, 'packing_algorithm': 'kk'}, id='dynamic-kk'),
+])
+def test_transformers_batch_context_loss_matches_independent_micro_batch_computation(factory, options):
+    rows = data()
+    model, oracle = BatchContextHarness(factory()), BatchContextHarness(factory())
+    kwargs = {}
+    if isinstance(model.optimizer_group['default'].loss_instance, GRPOLoss):
+        kwargs = {'advantages': [1., -.5, .7, -1.2, .4, -.3, .9, -.8],
+                  'old_logps': [[0.] * len(row['labels']) for row in rows]}
+
+    torch.manual_seed(7)
+    actual = model.run(rows, **options, **kwargs)
+
+    if options:
+        assert len(model.calls) > 1
+    else:
+        assert len(model.calls) == 1
+    assert sorted(index for indices, _ in model.calls for index in indices) == list(range(len(rows)))
+    assert [sync for _, sync in model.calls] == [False] * (len(model.calls) - 1) + [True]
+    assert model.optimizer_group['default'].cur_step == 1
+
+    # Replay the same groups directly through the loss: batch statistics must be
+    # recomputed for each group, rather than copied from the full update window.
+    torch.manual_seed(7)
+    terms = []
+    loss = oracle.optimizer_group['default'].loss_instance
+    for indices, _ in model.calls:
+        outputs = oracle.forward(inputs=[rows[index] for index in indices])
+        inputs = oracle.optimizer_group['default'].train_status.inputs
+        micro_kwargs = {key: [value[index] for index in indices] for key, value in kwargs.items()}
+        result = loss(inputs, outputs, **micro_kwargs)
+        scale = loss.micro_batch_scale(rows, indices)
+        terms.append(result['loss'] * scale)
+    expected = torch.stack(terms).sum()
+    expected.backward()
+
+    assert actual['loss'] == pytest.approx(expected.item())
+    assert torch.isfinite(model.weight.grad)
+    torch.testing.assert_close(model.weight.grad, oracle.weight.grad)
+
+
+@pytest.fixture
+def megatron_model(monkeypatch):
+    from twinkle.model.megatron.megatron import MegatronModel
+
+    # MCore is optional on CPU; stop after input batching and before scheduler setup.
+    core = ModuleType('megatron.core')
+    core.parallel_state = SimpleNamespace()
+    pipeline = ModuleType('megatron.core.pipeline_parallel')
+    pipeline.get_forward_backward_func = Mock(side_effect=AssertionError('scheduler must not run'))
+    monkeypatch.setitem(sys.modules, 'megatron.core', core)
+    monkeypatch.setitem(sys.modules, 'megatron.core.pipeline_parallel', pipeline)
+    processor = InputProcessor(framework='megatron')
+    processor.process_pipeline = [processor.prepare_inputs, processor.collate_fn]
+    group = BaseOptimizerGroup(loss_instance=ClipCovLoss())
+    group.processor = processor
+    chunk = SimpleNamespace(config=SimpleNamespace())
+    model = SimpleNamespace(
+        _lazy_wrap_model=lambda: None,
+        _lazy_finish_param_config=lambda: None,
+        _get_default_group=lambda: 'default',
+        _not_encoded=lambda _: False,
+        optimizer_group={'default': group},
+        strategy=SimpleNamespace(unwrap_model=lambda _: [chunk]),
+        model=[chunk],
+        variable_seq_lengths=False,
+        device_mesh=SimpleNamespace(cp_world_size=1),
+        _router_replay_enabled=False,
+        _maybe_apply_router_replay=lambda: None,
+    )
+    return model, unwrap(MegatronModel.forward_backward)
+
+
+@pytest.mark.parametrize('factory', BATCH_CONTEXT_LOSS_FACTORIES)
+@pytest.mark.parametrize('forward_only', [False, True])
+@pytest.mark.parametrize('input_count', [2, 4])
+def test_megatron_batch_context_losses_reach_schedule_setup(
+    megatron_model, monkeypatch, factory, forward_only, input_count,
+):
+    from twinkle.model.megatron import megatron
+
+    class ScheduleSetupReached(Exception):
+        pass
+
+    # Reaching this point proves both one and multiple micro-batches are accepted.
+    setup = Mock(side_effect=ScheduleSetupReached)
+    monkeypatch.setattr(megatron, 'RouterReplay', setup)
+    model, forward_backward = megatron_model
+    model.optimizer_group['default'].loss_instance = factory()
+    model.optimizer_group['default'].gradient_accumulation_steps = 4
+    rows = [{'input_ids': [1, 2], 'labels': [1, 2]} for _ in range(input_count)]
+
+    with pytest.raises(ScheduleSetupReached):
+        forward_backward(model, inputs=rows, micro_batch_size=2, forward_only=forward_only)
+
+    setup.assert_called_once()

tests/model/test_ppo_micro_batch.py

@@ -0,0 +1,74 @@
+"""PPO actor and critic match the full local batch under fixed and dynamic splitting."""
+import pytest
+import torch
+
+from twinkle.loss import PPOLoss, PPOValueLoss
+from .test_grpo_micro_batch import data
+from .test_recipe_micro_batch import Harness
+
+
+class PPOHarness(Harness):
+    def forward(self, *, inputs, **kwargs):
+        outputs = super().forward(inputs=inputs, **kwargs)
+        status = self.optimizer_group['default'].train_status
+        width = status.inputs['labels'].shape[1]
+        if any('completion_mask' in row for row in inputs):
+            status.inputs['completion_mask'] = torch.tensor([
+                row['completion_mask'] + [0] * (width - len(row['completion_mask'])) for row in inputs
+            ])
+        return outputs
+
+
+@pytest.mark.parametrize('factory', [
+    pytest.param(PPOLoss, id='policy-token-mean'),
+    pytest.param(lambda: PPOLoss(loss_agg_mode='seq-mean-token-mean'), id='policy-sequence-mean'),
+    pytest.param(PPOValueLoss, id='value-token-mean'),
+])
+@pytest.mark.parametrize('options', [
+    pytest.param({'micro_batch_size': 2}, id='fixed'),
+    pytest.param({'micro_batch_size': 4}, id='fixed-short-tail'),
+    pytest.param({'micro_batch_size': 4, 'dynamic_batching': True,
+                  'max_tokens_per_micro_batch': 18, 'packing_algorithm': 'ffd'}, id='dynamic-ffd'),
+    pytest.param({'micro_batch_size': 4, 'dynamic_batching': True,
+                  'max_tokens_per_micro_batch': 18, 'packing_algorithm': 'kk'}, id='dynamic-kk'),
+])
+@pytest.mark.parametrize('completion_mask', [False, True])
+def test_ppo_loss_and_gradients_match_unsplit_batch(factory, options, completion_mask):
+    baseline, internal = PPOHarness(factory()), PPOHarness(factory())
+    rows = data()[:6]
+    for index, row in enumerate(rows):
+        row['labels'][0] = -100
+        if completion_mask:
+            row['completion_mask'] = [int((position + index) % 3 != 0) for position in range(len(row['labels']))]
+    kwargs = {
+        'advantages': [1., -.5, .7, -1.2, .4, -.3],
+        'old_logps': [[0.] * len(row['labels']) for row in rows],
+        'old_values': [[0.] * len(row['labels']) for row in rows],
+        'returns': [[.2] * len(row['labels']) for row in rows],
+    }
+
+    expected = baseline.run(rows, **kwargs)
+    actual = internal.run(rows, **options, **kwargs)
+
+    assert actual['loss'] == pytest.approx(expected['loss'])
+    torch.testing.assert_close(internal.normalized_gradient(), baseline.normalized_gradient())
+
+
+@pytest.mark.parametrize('factory', [PPOLoss, PPOValueLoss])
+def test_empty_ppo_micro_batch_keeps_loss_and_gradient_finite(factory):
+    baseline, internal = PPOHarness(factory()), PPOHarness(factory())
+    rows = data()[:4]
+    for row in rows[:2]:
+        row['labels'] = [-100] * len(row['labels'])
+    kwargs = {
+        'advantages': [1., -.5, .7, -1.2],
+        'old_logps': [[0.] * len(row['labels']) for row in rows],
+        'old_values': [[0.] * len(row['labels']) for row in rows],
+        'returns': [[.2] * len(row['labels']) for row in rows],
+    }
+
+    expected = baseline.run(rows, **kwargs)
+    actual = internal.run(rows, micro_batch_size=2, **kwargs)
+
+    assert actual['loss'] == pytest.approx(expected['loss'])
+    torch.testing.assert_close(internal.normalized_gradient(), baseline.normalized_gradient())

tests/model/test_recipe_micro_batch.py

@@ -13,7 +13,6 @@
 from twinkle.loss.mopd import MOPDLoss
 from twinkle.loss.reward import RewardLoss
 from twinkle.loss.seq_cls import SeqClsLoss
-from twinkle.loss.value import PPOValueLoss
 from .test_grpo_micro_batch import TrainingHarness, data
 
 
@@ -50,7 +49,7 @@ def test_sample_mean_losses_match_whole_batch_with_a_short_micro_batch(factory):
     torch.testing.assert_close(internal.normalized_gradient(), baseline.normalized_gradient())
 
 
-@pytest.mark.parametrize("factory", [PPOValueLoss, GKDLoss, LigerFusedLinearCrossEntropyLoss])
+@pytest.mark.parametrize("factory", [GKDLoss, LigerFusedLinearCrossEntropyLoss])
 def test_token_losses_match_explicit_fixed_micro_batches(factory):
     baseline, internal = Harness(factory()), Harness(factory())
     rows = data()[:6]