diff --git a/slime/rollout/on_policy_distillation.py b/slime/rollout/on_policy_distillation.py index 919097434..b866ace08 100644 --- a/slime/rollout/on_policy_distillation.py +++ b/slime/rollout/on_policy_distillation.py @@ -50,8 +50,11 @@ def post_process_rewards(args, samples: list[Sample], **kwargs): torch.tensor([item[0] for item in reward["meta_info"]["input_token_logprobs"][1:]], dtype=torch.float32) for reward in raw_rewards ] + # NB: t_log_prob[-0:] is the whole tensor, not an empty one, so a zero-length + # response (e.g. an immediate EOS) must be handled explicitly to keep + # teacher_log_probs aligned with response_length. teacher_log_probs = [ - t_log_prob[-response_length:] + t_log_prob[-response_length:] if response_length > 0 else t_log_prob[:0] for t_log_prob, response_length in zip(teacher_log_probs, response_lengths, strict=False) ] diff --git a/slime/rollout/sft_rollout.py b/slime/rollout/sft_rollout.py index 4407463d2..8754750f6 100644 --- a/slime/rollout/sft_rollout.py +++ b/slime/rollout/sft_rollout.py @@ -57,7 +57,10 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False): sample.tokens = token_ids sample.response_length = response_length sample.reward = 0 - sample.loss_mask = loss_mask[-response_length:] + # NB: loss_mask[-0:] is the whole list, not an empty one, so a sample with no + # trainable tokens (response_length == 0) must be handled explicitly. Otherwise + # the downstream invariant len(loss_mask) == response_length is violated. + sample.loss_mask = loss_mask[-response_length:] if response_length > 0 else [] if i == 0 and not SAMPLE_PRINTED: logger.info( diff --git a/tests/test_zero_length_response.py b/tests/test_zero_length_response.py new file mode 100644 index 000000000..6e707cbb1 --- /dev/null +++ b/tests/test_zero_length_response.py @@ -0,0 +1,127 @@ +"""CPU unit tests for zero-length response handling in rollout post-processing. + +Pins the contract that a response of length 0 must produce an *empty* +loss_mask / teacher_log_probs, not the whole sequence. Python's slicing +semantics make ``seq[-0:]`` return the entire sequence, so every +``seq[-response_length:]`` in the rollout path needs an explicit guard +for ``response_length == 0``: + + * ``slime/rollout/sft_rollout.py`` — a fully-masked SFT sample (no + trainable tokens) used to end up with ``loss_mask`` covering the + whole prompt, violating the downstream invariant + ``len(loss_mask) == response_length`` (asserted in + ``slime/ray/rollout.py`` and ``slime/utils/types.py``) and crashing + the rollout. + + * ``slime/rollout/on_policy_distillation.py`` — an empty generation + (e.g. immediate EOS) used to keep the full-length teacher log-prob + tensor instead of an empty one, misaligning the OPD KL computation. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import slime.rollout.sft_rollout as sft_rollout +from slime.rollout.on_policy_distillation import post_process_rewards +from slime.utils.mask_utils import get_response_lengths +from slime.utils.types import Sample + +NUM_GPUS = 0 + + +class _FakeMaskGenerator: + """Stands in for MultiTurnLossMaskGenerator without needing a tokenizer.""" + + def __init__(self, token_ids, loss_mask): + self._token_ids = token_ids + self._loss_mask = loss_mask + + def get_loss_mask(self, messages, tools=None): + return self._token_ids, self._loss_mask + + def get_response_lengths(self, loss_masks): + return get_response_lengths(loss_masks) + + +class _FakeDataBuffer: + def __init__(self, samples): + self._samples = samples + + def get_samples(self, n): + return self._samples[:n] + + +def _run_sft_rollout(monkeypatch, token_ids, loss_mask): + monkeypatch.setattr(sft_rollout, "TOKENIZER", object()) + monkeypatch.setattr(sft_rollout, "PROCESSOR", object()) + monkeypatch.setattr(sft_rollout, "MASK_GENERATOR", _FakeMaskGenerator(token_ids, loss_mask)) + monkeypatch.setattr(sft_rollout, "SAMPLE_PRINTED", True) + + args = SimpleNamespace( + hf_checkpoint=None, + loss_mask_type=None, + rollout_global_dataset=True, + rollout_batch_size=1, + ) + data_buffer = _FakeDataBuffer([(Sample(prompt=[{"role": "user", "content": "hi"}]),)]) + (sample,) = sft_rollout.generate_rollout(args, 0, data_buffer)[0] + return sample + + +@pytest.mark.unit +def test_sft_rollout_fully_masked_sample_yields_empty_loss_mask(monkeypatch): + # No assistant turn / template mismatch -> no trainable tokens at all. + sample = _run_sft_rollout(monkeypatch, [10, 11, 12, 13], [0, 0, 0, 0]) + + assert sample.response_length == 0 + # Must be empty; the buggy `loss_mask[-0:]` returned the whole mask here. + assert sample.loss_mask == [] + assert len(sample.loss_mask) == sample.response_length + + +@pytest.mark.unit +def test_sft_rollout_normal_sample_keeps_tail_loss_mask(monkeypatch): + sample = _run_sft_rollout(monkeypatch, [10, 11, 12, 13], [0, 0, 1, 1]) + + assert sample.response_length == 2 + assert sample.loss_mask == [1, 1] + assert len(sample.loss_mask) == sample.response_length + + +class _StubOPDSample: + """Minimal stand-in for Sample as used by post_process_rewards.""" + + def __init__(self, reward_payload, response_length): + self._reward_payload = reward_payload + self.response_length = response_length + self.teacher_log_probs = None + + def get_reward_value(self, args): + return self._reward_payload + + +def _opd_reward_payload(logprobs): + # Mirrors sglang's return_logprob format: the first entry has no logprob. + return {"meta_info": {"input_token_logprobs": [(None, 0)] + [(lp, i + 1) for i, lp in enumerate(logprobs)]}} + + +@pytest.mark.unit +def test_opd_zero_length_response_yields_empty_teacher_log_probs(): + sample = _StubOPDSample(_opd_reward_payload([-0.1, -0.2, -0.3]), response_length=0) + + post_process_rewards(None, [sample]) + + assert sample.teacher_log_probs.shape == (0,) + + +@pytest.mark.unit +def test_opd_normal_response_keeps_tail_teacher_log_probs(): + sample = _StubOPDSample(_opd_reward_payload([-0.1, -0.2, -0.3]), response_length=2) + + post_process_rewards(None, [sample]) + + assert torch.allclose(sample.teacher_log_probs, torch.tensor([-0.2, -0.3]))