From 78a8c93236dc7d8b4643858260f4d4a65186bafa Mon Sep 17 00:00:00 2001
From: meichangsu1 <1484603386@qq.com>
Date: Fri, 30 Oct 2026 17:29:28 +0800
Subject: [PATCH 1/2] fix loss  micro_batch_scale

---
 src/twinkle/loss/base.py                      |  12 ++
 src/twinkle/loss/clip_cov.py                  |   2 +
 src/twinkle/loss/grpo.py                      |  16 ++
 src/twinkle/loss/infonce.py                   |   5 +
 src/twinkle/loss/kl_cov.py                    |   2 +
 src/twinkle/loss/utils.py                     |  13 +-
 src/twinkle/loss/value.py                     |   5 +-
 src/twinkle/model/megatron/megatron.py        |   2 +
 .../model/transformers/transformers.py        |   1 +
 tests/loss/test_ppo.py                        |  25 +++
 tests/model/test_batch_context_loss.py        | 156 ++++++++++++++++++
 tests/model/test_ppo_micro_batch.py           |  74 +++++++++
 tests/model/test_recipe_micro_batch.py        |   3 +-
 13 files changed, 310 insertions(+), 6 deletions(-)
 create mode 100644 tests/model/test_batch_context_loss.py
 create mode 100644 tests/model/test_ppo_micro_batch.py

diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py
index 4e017ccbe..a7d140ef0 100644
--- a/src/twinkle/loss/base.py
+++ b/src/twinkle/loss/base.py
@@ -8,10 +8,22 @@ class Loss:
     require_entropy = False
     require_logps = True
     require_values = False
+    # This loss needs all local inputs in one evaluation to preserve its batch context.
+    requires_full_batch = False
 
     def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput:
         ...
 
+    def validate_micro_batching(self, num_micro_batches: int) -> None:
+        """Validate the actual local plan; one micro-batch preserves the full batch context."""
+        if num_micro_batches <= 1 or not self.requires_full_batch:
+            return
+        raise NotImplementedError(
+            f'{type(self).__name__} requires the full local batch for its loss computation, but the '
+            f'current batch would be split into {num_micro_batches} micro-batches. This loss does not '
+            'support micro-batching, including dynamic batching. Increase micro_batch_size or the '
+            'dynamic token budget so the full local batch fits in a single micro-batch.')
+
     def micro_batch_scale(self, inputs: list[InputFeature], indices: list[int]) -> float:
         if len(indices) == len(inputs):
             return 1.0
diff --git a/src/twinkle/loss/clip_cov.py b/src/twinkle/loss/clip_cov.py
index 7833aa88d..9de7936e5 100644
--- a/src/twinkle/loss/clip_cov.py
+++ b/src/twinkle/loss/clip_cov.py
@@ -27,6 +27,8 @@ class ClipCovLoss(GRPOLoss):
         clip_cov_ub: Upper bound for covariance selection window (default 5.0).
     """
 
+    requires_full_batch = True
+
     def __init__(self, clip_cov_ratio: float = 0.0002, clip_cov_lb: float = 1.0, clip_cov_ub: float = 5.0, **kwargs):
         super().__init__(**kwargs)
         if clip_cov_ratio <= 0:
diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py
index 016e82988..a5799d3b8 100644
--- a/src/twinkle/loss/grpo.py
+++ b/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:
@@ -116,6 +117,10 @@ def __init__(
         self.require_entropy = entropy_coef > 0.0 or log_entropy or top_entropy_quantile < 1.0
         self.require_logits = self._distill_wrap
 
+    @property
+    def requires_full_batch(self):
+        return self.top_entropy_quantile < 1.0
+
     def micro_batch_scale(self, inputs, indices):
         return len(indices) / len(inputs)
 
@@ -722,6 +727,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',
diff --git a/src/twinkle/loss/infonce.py b/src/twinkle/loss/infonce.py
index 2efff16a4..9e51d8b60 100644
--- a/src/twinkle/loss/infonce.py
+++ b/src/twinkle/loss/infonce.py
@@ -215,6 +215,10 @@ def __init__(
         self.include_dd = include_dd
         self.process_group = process_group
 
+    @property
+    def requires_full_batch(self):
+        return self.use_batch
+
     def _gather_across_dp(self, sentences: torch.Tensor, labels: torch.Tensor):
         """All-gather embeddings & labels across DP ranks; only local shard keeps grad.
 
@@ -464,6 +468,7 @@ class OnlineContrastiveLoss(EmbeddingLoss):
     require_logits = True
     require_entropy = False
     require_logps = False
+    requires_full_batch = True
 
     def __init__(self, margin: float = 0.5, distance_metric='cosine', mrl_dims=None, **kwargs):
         super().__init__(mrl_dims=mrl_dims)
diff --git a/src/twinkle/loss/kl_cov.py b/src/twinkle/loss/kl_cov.py
index bb89d481d..0710979b9 100644
--- a/src/twinkle/loss/kl_cov.py
+++ b/src/twinkle/loss/kl_cov.py
@@ -30,6 +30,8 @@ class KLCovLoss(GRPOLoss):
         ppo_kl_coef: Coefficient for the KL penalty term (default 1.0).
     """
 
+    requires_full_batch = True
+
     def __init__(self, kl_cov_ratio: float = 0.0002, ppo_kl_coef: float = 1.0, **kwargs):
         super().__init__(**kwargs)
         if kl_cov_ratio <= 0:
diff --git a/src/twinkle/loss/utils.py b/src/twinkle/loss/utils.py
index 9d1cad523..11452f9ac 100644
--- a/src/twinkle/loss/utils.py
+++ b/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)
diff --git a/src/twinkle/loss/value.py b/src/twinkle/loss/value.py
index 63f1f5851..5e3367c6c 100644
--- a/src/twinkle/loss/value.py
+++ b/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,
diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py
index 068b5fc82..e343d3f10 100644
--- a/src/twinkle/model/megatron/megatron.py
+++ b/src/twinkle/model/megatron/megatron.py
@@ -503,6 +503,8 @@ def forward_backward(self,
             variable_seq_lengths=self.variable_seq_lengths,
             attention_mask_type=getattr(unwrapped_model.config, 'attention_mask_type', None),
         )
+        if compute_loss:
+            loss_instance.validate_micro_batching(len(inputs))
 
         # Get parallelism settings for sequence padding and splitting
         cp_size = self.device_mesh.cp_world_size
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 315a69f50..20909c565 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -1540,6 +1540,7 @@ def _forward_backward_micro_batches(
         local_batch_size = len(inputs)
         plan = self._build_micro_batch_plan(inputs, config, optimizer_config)
         loss_instance = optimizer_config.get_loss(loss_key)
+        loss_instance.validate_micro_batching(len(plan))
         total_loss = 0.0
         outputs = {}
         for micro_batch_index, indices in enumerate(plan):
diff --git a/tests/loss/test_ppo.py b/tests/loss/test_ppo.py
index 882ea5797..921e4ad39 100644
--- a/tests/loss/test_ppo.py
+++ b/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)(
diff --git a/tests/model/test_batch_context_loss.py b/tests/model/test_batch_context_loss.py
new file mode 100644
index 000000000..2a475328b
--- /dev/null
+++ b/tests/model/test_batch_context_loss.py
@@ -0,0 +1,156 @@
+"""Batch-dependent objectives must fail before a micro-batched loss is evaluated."""
+import sys
+from inspect import unwrap
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from twinkle.loss import CrossEntropyLoss, GRPOLoss
+from twinkle.loss.clip_cov import ClipCovLoss
+from twinkle.loss.grpo import BNPOLoss, CISPOLoss, PPOLoss
+from twinkle.loss.infonce import 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_DEPENDENT_LOSSES = [
+    pytest.param(InfonceLoss(), id='infonce-batch-negatives'),
+    pytest.param(OnlineContrastiveLoss(), id='online-contrastive'),
+    pytest.param(ClipCovLoss(), id='clip-cov'),
+    pytest.param(KLCovLoss(), id='kl-cov'),
+    pytest.param(GRPOLoss(top_entropy_quantile=.5), id='grpo-entropy-filter'),
+    pytest.param(PPOLoss(top_entropy_quantile=.5), id='ppo-entropy-filter'),
+    pytest.param(CISPOLoss(top_entropy_quantile=.5), id='cispo-entropy-filter'),
+    pytest.param(BNPOLoss(top_entropy_quantile=.5, token_mean_scope='global'), id='bnpo-entropy-filter'),
+]
+
+
+@pytest.mark.parametrize('loss', BATCH_DEPENDENT_LOSSES)
+@pytest.mark.parametrize('algorithm', [None, 'ffd', 'kk'])
+def test_transformers_rejects_batch_dependent_split_before_forward(loss, algorithm):
+    model = TrainingHarness(loss)
+    options = {'micro_batch_size': 2}
+    if algorithm is not None:
+        options.update(dynamic_batching=True, max_tokens_per_micro_batch=18, packing_algorithm=algorithm)
+
+    with pytest.raises(NotImplementedError, match=f'{type(loss).__name__}.*full local batch') as error:
+        model.run(data(), **options)
+
+    assert 'would be split into' in str(error.value)
+    assert model.calls == []
+    assert model.weight.grad is None
+    assert model.optimizer_group['default'].cur_step == 0
+
+
+@pytest.mark.parametrize('loss', BATCH_DEPENDENT_LOSSES)
+@pytest.mark.parametrize('algorithm', [None, 'ffd', 'kk'])
+def test_transformers_allows_batch_dependent_loss_when_plan_does_not_split(loss, algorithm):
+    model = TrainingHarness(loss)
+    model._forward_backward_micro_batch = Mock(return_value={'loss': 0.})
+    rows = data()
+    options = {'micro_batch_size': len(rows)}
+    if algorithm is not None:
+        options.update(dynamic_batching=True, max_tokens_per_micro_batch=100, packing_algorithm=algorithm)
+
+    model.run(rows, **options)
+
+    model._forward_backward_micro_batch.assert_called_once()
+    assert len(model._forward_backward_micro_batch.call_args.kwargs['inputs']) == len(rows)
+
+
+@pytest.mark.parametrize('loss', BATCH_DEPENDENT_LOSSES)
+@pytest.mark.parametrize('options', [
+    pytest.param({}, id='unsplit'),
+    pytest.param({'micro_batch_size': 4}, id='fixed-short-batch'),
+    pytest.param({'micro_batch_size': 4, 'dynamic_batching': True,
+                  'max_tokens_per_micro_batch': 100, 'packing_algorithm': 'ffd'}, id='dynamic-ffd'),
+    pytest.param({'micro_batch_size': 4, 'dynamic_batching': True,
+                  'max_tokens_per_micro_batch': 100, 'packing_algorithm': 'kk'}, id='dynamic-kk'),
+])
+def test_transformers_allows_one_actual_micro_batch_despite_configured_ga(loss, options):
+    model = TrainingHarness(loss)
+    model.optimizer_group['default'].gradient_accumulation_steps = 4
+    model._forward_backward_micro_batch = Mock(return_value={'loss': 0.})
+    rows = data()[:2]
+
+    model.run(rows, **options)
+
+    model._forward_backward_micro_batch.assert_called_once()
+    actual = model._forward_backward_micro_batch.call_args.kwargs['inputs']
+    assert sorted(actual, key=lambda row: row['input_ids'][0]) == rows
+
+
+@pytest.mark.parametrize('loss', [CrossEntropyLoss(), GRPOLoss(), InfonceLoss(use_batch=False)])
+def test_losses_without_batch_context_still_allow_splitting(loss):
+    loss.validate_micro_batching(4)
+
+
+@pytest.fixture
+def megatron_model(monkeypatch):
+    from twinkle.model.megatron.megatron import MegatronModel
+
+    # MCore is optional on CPU; no scheduler/model forward may run before the guard.
+    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('forward_only', [False, True])
+def test_megatron_rejects_batch_dependent_training_and_eval_split(megatron_model, forward_only):
+    model, forward_backward = megatron_model
+    rows = [{'input_ids': [1, 2], 'labels': [1, 2]} for _ in range(4)]
+
+    with pytest.raises(NotImplementedError, match='ClipCovLoss.*full local batch'):
+        forward_backward(model, inputs=rows, micro_batch_size=2, forward_only=forward_only)
+
+
+@pytest.mark.parametrize('forward_only', [False, True])
+@pytest.mark.parametrize('micro_batch_size', [None, 2, 4])
+def test_megatron_allows_one_actual_micro_batch(megatron_model, monkeypatch, forward_only, micro_batch_size):
+    from twinkle.model.megatron import megatron
+
+    class ScheduleSetupReached(Exception):
+        pass
+
+    # Stop after the guard, before reaching the optional MCore scheduler or GPU operations.
+    setup = Mock(side_effect=ScheduleSetupReached)
+    monkeypatch.setattr(megatron, 'RouterReplay', setup)
+    model, forward_backward = megatron_model
+    model.optimizer_group['default'].gradient_accumulation_steps = 4
+    rows = [{'input_ids': [1, 2], 'labels': [1, 2]} for _ in range(2)]
+
+    with pytest.raises(ScheduleSetupReached):
+        forward_backward(model, inputs=rows, micro_batch_size=micro_batch_size, forward_only=forward_only)
+
+    setup.assert_called_once()
+
+
+def test_batch_dependent_error_reports_actual_count_in_english():
+    with pytest.raises(NotImplementedError, match='split into 3 micro-batches') as error:
+        ClipCovLoss().validate_micro_batching(3)
+
+    assert str(error.value).isascii()
diff --git a/tests/model/test_ppo_micro_batch.py b/tests/model/test_ppo_micro_batch.py
new file mode 100644
index 000000000..ff56836e5
--- /dev/null
+++ b/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())
diff --git a/tests/model/test_recipe_micro_batch.py b/tests/model/test_recipe_micro_batch.py
index 32a00ffc5..e7ad27bb5 100644
--- a/tests/model/test_recipe_micro_batch.py
+++ b/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]

From 258f85845ae0c42c6529ab96a9bd87e9316aa84a Mon Sep 17 00:00:00 2001
From: meichangsu1 <1484603386@qq.com>
Date: Mon, 2 Nov 2026 15:36:35 +0800
Subject: [PATCH 2/2] fix microbatch

---
 src/twinkle/loss/base.py                      |  12 --
 src/twinkle/loss/clip_cov.py                  |   2 -
 src/twinkle/loss/grpo.py                      |   5 +-
 src/twinkle/loss/infonce.py                   |   8 +-
 src/twinkle/loss/kl_cov.py                    |   2 -
 src/twinkle/model/megatron/megatron.py        |   3 -
 .../model/transformers/transformers.py        |   1 -
 tests/model/test_batch_context_loss.py        | 176 +++++++++---------
 8 files changed, 93 insertions(+), 116 deletions(-)

diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py
index 7a665d353..bc02ebe92 100644
--- a/src/twinkle/loss/base.py
+++ b/src/twinkle/loss/base.py
@@ -8,22 +8,10 @@ class Loss:
     require_entropy = False
     require_logps = True
     require_values = False
-    # This loss needs all local inputs in one evaluation to preserve its batch context.
-    requires_full_batch = False
 
     def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput:
         ...
 
-    def validate_micro_batching(self, num_micro_batches: int) -> None:
-        """Validate the actual local plan; one micro-batch preserves the full batch context."""
-        if num_micro_batches <= 1 or not self.requires_full_batch:
-            return
-        raise NotImplementedError(
-            f'{type(self).__name__} requires the full local batch for its loss computation, but the '
-            f'current batch would be split into {num_micro_batches} micro-batches. This loss does not '
-            'support micro-batching, including dynamic batching. Increase micro_batch_size or the '
-            'dynamic token budget so the full local batch fits in a single micro-batch.')
-
     def validate_inputs(self, inputs, **kwargs) -> None:
         """Reject a batch before any rank starts its forward pass.
 
diff --git a/src/twinkle/loss/clip_cov.py b/src/twinkle/loss/clip_cov.py
index 9de7936e5..7833aa88d 100644
--- a/src/twinkle/loss/clip_cov.py
+++ b/src/twinkle/loss/clip_cov.py
@@ -27,8 +27,6 @@ class ClipCovLoss(GRPOLoss):
         clip_cov_ub: Upper bound for covariance selection window (default 5.0).
     """
 
-    requires_full_batch = True
-
     def __init__(self, clip_cov_ratio: float = 0.0002, clip_cov_lb: float = 1.0, clip_cov_ub: float = 5.0, **kwargs):
         super().__init__(**kwargs)
         if clip_cov_ratio <= 0:
diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py
index 849871a4c..130edb03f 100644
--- a/src/twinkle/loss/grpo.py
+++ b/src/twinkle/loss/grpo.py
@@ -117,10 +117,6 @@ def __init__(
         self.require_entropy = entropy_coef > 0.0 or log_entropy or top_entropy_quantile < 1.0
         self.require_logits = self._distill_wrap
 
-    @property
-    def requires_full_batch(self):
-        return self.top_entropy_quantile < 1.0
-
     def micro_batch_scale(self, inputs, indices):
         return len(indices) / len(inputs)
 
@@ -637,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)
diff --git a/src/twinkle/loss/infonce.py b/src/twinkle/loss/infonce.py
index 9e51d8b60..928a36766 100644
--- a/src/twinkle/loss/infonce.py
+++ b/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.
@@ -215,10 +217,6 @@ def __init__(
         self.include_dd = include_dd
         self.process_group = process_group
 
-    @property
-    def requires_full_batch(self):
-        return self.use_batch
-
     def _gather_across_dp(self, sentences: torch.Tensor, labels: torch.Tensor):
         """All-gather embeddings & labels across DP ranks; only local shard keeps grad.
 
@@ -458,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.
@@ -468,7 +467,6 @@ class OnlineContrastiveLoss(EmbeddingLoss):
     require_logits = True
     require_entropy = False
     require_logps = False
-    requires_full_batch = True
 
     def __init__(self, margin: float = 0.5, distance_metric='cosine', mrl_dims=None, **kwargs):
         super().__init__(mrl_dims=mrl_dims)
diff --git a/src/twinkle/loss/kl_cov.py b/src/twinkle/loss/kl_cov.py
index 0710979b9..bb89d481d 100644
--- a/src/twinkle/loss/kl_cov.py
+++ b/src/twinkle/loss/kl_cov.py
@@ -30,8 +30,6 @@ class KLCovLoss(GRPOLoss):
         ppo_kl_coef: Coefficient for the KL penalty term (default 1.0).
     """
 
-    requires_full_batch = True
-
     def __init__(self, kl_cov_ratio: float = 0.0002, ppo_kl_coef: float = 1.0, **kwargs):
         super().__init__(**kwargs)
         if kl_cov_ratio <= 0:
diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py
index 804461a8a..77729ae3d 100644
--- a/src/twinkle/model/megatron/megatron.py
+++ b/src/twinkle/model/megatron/megatron.py
@@ -509,9 +509,6 @@ def forward_backward(self,
             variable_seq_lengths=self.variable_seq_lengths,
             attention_mask_type=attention_mask_type,
         )
-        if compute_loss:
-            loss_instance.validate_micro_batching(len(inputs))
-
         # 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);
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 815a0b9b5..3b527678d 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -1521,7 +1521,6 @@ def _forward_backward_micro_batches(
         local_batch_size = len(inputs)
         plan = self._build_micro_batch_plan(inputs, config, optimizer_config)
         loss_instance = optimizer_config.get_loss(loss_key)
-        loss_instance.validate_micro_batching(len(plan))
         total_loss = 0.0
         outputs = {}
         for micro_batch_index, indices in enumerate(plan):
diff --git a/tests/model/test_batch_context_loss.py b/tests/model/test_batch_context_loss.py
index 2a475328b..206a413c9 100644
--- a/tests/model/test_batch_context_loss.py
+++ b/tests/model/test_batch_context_loss.py
@@ -1,97 +1,111 @@
-"""Batch-dependent objectives must fail before a micro-batched loss is evaluated."""
+"""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 CrossEntropyLoss, GRPOLoss
+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 InfonceLoss, OnlineContrastiveLoss
+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_DEPENDENT_LOSSES = [
-    pytest.param(InfonceLoss(), id='infonce-batch-negatives'),
-    pytest.param(OnlineContrastiveLoss(), id='online-contrastive'),
-    pytest.param(ClipCovLoss(), id='clip-cov'),
-    pytest.param(KLCovLoss(), id='kl-cov'),
-    pytest.param(GRPOLoss(top_entropy_quantile=.5), id='grpo-entropy-filter'),
-    pytest.param(PPOLoss(top_entropy_quantile=.5), id='ppo-entropy-filter'),
-    pytest.param(CISPOLoss(top_entropy_quantile=.5), id='cispo-entropy-filter'),
-    pytest.param(BNPOLoss(top_entropy_quantile=.5, token_mean_scope='global'), id='bnpo-entropy-filter'),
+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'),
 ]
 
 
-@pytest.mark.parametrize('loss', BATCH_DEPENDENT_LOSSES)
-@pytest.mark.parametrize('algorithm', [None, 'ffd', 'kk'])
-def test_transformers_rejects_batch_dependent_split_before_forward(loss, algorithm):
-    model = TrainingHarness(loss)
-    options = {'micro_batch_size': 2}
-    if algorithm is not None:
-        options.update(dynamic_batching=True, max_tokens_per_micro_batch=18, packing_algorithm=algorithm)
-
-    with pytest.raises(NotImplementedError, match=f'{type(loss).__name__}.*full local batch') as error:
-        model.run(data(), **options)
-
-    assert 'would be split into' in str(error.value)
-    assert model.calls == []
-    assert model.weight.grad is None
-    assert model.optimizer_group['default'].cur_step == 0
-
-
-@pytest.mark.parametrize('loss', BATCH_DEPENDENT_LOSSES)
-@pytest.mark.parametrize('algorithm', [None, 'ffd', 'kk'])
-def test_transformers_allows_batch_dependent_loss_when_plan_does_not_split(loss, algorithm):
-    model = TrainingHarness(loss)
-    model._forward_backward_micro_batch = Mock(return_value={'loss': 0.})
-    rows = data()
-    options = {'micro_batch_size': len(rows)}
-    if algorithm is not None:
-        options.update(dynamic_batching=True, max_tokens_per_micro_batch=100, packing_algorithm=algorithm)
-
-    model.run(rows, **options)
-
-    model._forward_backward_micro_batch.assert_called_once()
-    assert len(model._forward_backward_micro_batch.call_args.kwargs['inputs']) == len(rows)
-
-
-@pytest.mark.parametrize('loss', BATCH_DEPENDENT_LOSSES)
+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-short-batch'),
+    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': 100, 'packing_algorithm': 'ffd'}, id='dynamic-ffd'),
+                  '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': 100, 'packing_algorithm': 'kk'}, id='dynamic-kk'),
+                  'max_tokens_per_micro_batch': 18, 'packing_algorithm': 'kk'}, id='dynamic-kk'),
 ])
-def test_transformers_allows_one_actual_micro_batch_despite_configured_ga(loss, options):
-    model = TrainingHarness(loss)
-    model.optimizer_group['default'].gradient_accumulation_steps = 4
-    model._forward_backward_micro_batch = Mock(return_value={'loss': 0.})
-    rows = data()[:2]
-
-    model.run(rows, **options)
-
-    model._forward_backward_micro_batch.assert_called_once()
-    actual = model._forward_backward_micro_batch.call_args.kwargs['inputs']
-    assert sorted(actual, key=lambda row: row['input_ids'][0]) == rows
-
-
-@pytest.mark.parametrize('loss', [CrossEntropyLoss(), GRPOLoss(), InfonceLoss(use_batch=False)])
-def test_losses_without_batch_context_still_allow_splitting(loss):
-    loss.validate_micro_batching(4)
+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; no scheduler/model forward may run before the guard.
+    # 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')
@@ -119,38 +133,26 @@ def megatron_model(monkeypatch):
     return model, unwrap(MegatronModel.forward_backward)
 
 
+@pytest.mark.parametrize('factory', BATCH_CONTEXT_LOSS_FACTORIES)
 @pytest.mark.parametrize('forward_only', [False, True])
-def test_megatron_rejects_batch_dependent_training_and_eval_split(megatron_model, forward_only):
-    model, forward_backward = megatron_model
-    rows = [{'input_ids': [1, 2], 'labels': [1, 2]} for _ in range(4)]
-
-    with pytest.raises(NotImplementedError, match='ClipCovLoss.*full local batch'):
-        forward_backward(model, inputs=rows, micro_batch_size=2, forward_only=forward_only)
-
-
-@pytest.mark.parametrize('forward_only', [False, True])
-@pytest.mark.parametrize('micro_batch_size', [None, 2, 4])
-def test_megatron_allows_one_actual_micro_batch(megatron_model, monkeypatch, forward_only, micro_batch_size):
+@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
 
-    # Stop after the guard, before reaching the optional MCore scheduler or GPU operations.
+    # 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(2)]
+    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=micro_batch_size, forward_only=forward_only)
+        forward_backward(model, inputs=rows, micro_batch_size=2, forward_only=forward_only)
 
     setup.assert_called_once()
-
-
-def test_batch_dependent_error_reports_actual_count_in_english():
-    with pytest.raises(NotImplementedError, match='split into 3 micro-batches') as error:
-        ClipCovLoss().validate_micro_batching(3)
-
-    assert str(error.value).isascii()