diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 7a0a84378..130edb03f 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: @@ -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', diff --git a/src/twinkle/loss/infonce.py b/src/twinkle/loss/infonce.py index 2efff16a4..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. @@ -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. 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 b69dfdb63..77729ae3d 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/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); 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..206a413c9 --- /dev/null +++ b/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() 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]