diff --git a/tests/experimental/test_minillm_trainer.py b/tests/experimental/test_minillm_trainer.py index 59c0be9f306..619e705dd46 100644 --- a/tests/experimental/test_minillm_trainer.py +++ b/tests/experimental/test_minillm_trainer.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import types + import pytest import torch from datasets import DatasetDict, load_dataset @@ -80,3 +82,64 @@ def test_init_with_eval_dataset(self, eval_dataset_type): assert set(trainer.eval_dataset.keys()) == {"data1", "data2"} else: assert trainer.eval_dataset is eval_dataset + + +class TestMiniLLMComputeAdvantage: + def test_length_normalization_ignores_padding(self): + """A constant reward of 1.0 with gamma=1.0 must give an advantage of exactly 1.0 at every valid position, + however long the batch is padded. Counting padded slots toward the length, as the old 1e-4 fill did, lowers the + first position of a short completion (0.9514 for a single valid token padded to 512), and a floor on the + discounted length would deflate late positions once gamma < 1.""" + for gamma in (1.0, 0.9): + stub = types.SimpleNamespace(gamma=gamma, length_normalization=True) + for length in (1, 128, 512): + mask = torch.zeros(1, 512) + mask[0, :length] = 1 + advantages = MiniLLMTrainer._compute_advantage(stub, torch.zeros(1, 512), torch.ones(1, 512), mask) + # With gamma < 1 the discounted length of a late position is tiny (0.9**300 is 1.9e-14), so any floor + # on the denominator would deflate it; the expected value is exactly the reward at every position. + torch.testing.assert_close(advantages[0, :length], torch.ones(length)) + torch.testing.assert_close(advantages[0, length:], torch.zeros(512 - length)) + + advantages = MiniLLMTrainer._compute_advantage( + stub, torch.zeros(1, 8), torch.ones(1, 8), torch.zeros(1, 8) + ) + assert torch.isfinite(advantages).all() + torch.testing.assert_close(advantages, torch.zeros(1, 8)) + + +class TestMiniLLMComputeLossMask: + def test_advantage_mask_excludes_padded_completion_tokens(self, monkeypatch): + """The mask handed to `_compute_advantage` must come from the -100 fill in `labels`; built from `input_ids` + it was all ones, so padded slots kept a reward and leaked into every earlier advantage.""" + + class DummyModel: + def eval(self): + return self + + def __call__(self, input_ids, attention_mask, use_cache): + return types.SimpleNamespace(logits=torch.zeros(*input_ids.shape, 10)) + + trainer = object.__new__(MiniLLMTrainer) + trainer.teacher_model = DummyModel() + trainer.kd_temperature = 1.0 + trainer.rkl_advantage = True + + completion_mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]]) + inputs = { + "prompt_ids": torch.randint(1, 10, (2, 3)), + "prompt_mask": torch.ones(2, 3, dtype=torch.long), + "completion_ids": torch.randint(1, 10, (2, 4)), + "completion_mask": completion_mask, + } + + class AdvantageCaptured(Exception): + pass + + def spy(**kwargs): + assert torch.equal(kwargs["mask"], completion_mask.bool()) + raise AdvantageCaptured + + monkeypatch.setattr(trainer, "_compute_advantage", spy) + with pytest.raises(AdvantageCaptured): + trainer.compute_loss(DummyModel(), inputs) diff --git a/trl/experimental/minillm/minillm_trainer.py b/trl/experimental/minillm/minillm_trainer.py index 646f1355426..e37a8af5882 100644 --- a/trl/experimental/minillm/minillm_trainer.py +++ b/trl/experimental/minillm/minillm_trainer.py @@ -336,10 +336,12 @@ def _compute_advantage( advantages = advantages.flip(1).cumsum(dim=1).flip(1) if self.length_normalization: - mask = torch.where(mask < 0.5, 1e-4, mask) - lengths = mask * gamma_pow - lengths = lengths.flip(1).cumsum(dim=1).flip(1) - advantages = advantages / lengths + # Only unmasked positions count toward the discounted length, so the advantage of a completion does + # not depend on how much padding the batch carries. A length is zero only where every later position + # is masked, and there the numerator is zero too, so dividing by one keeps that tail at zero without + # touching valid positions whose discounted length is small. + lengths = (mask * gamma_pow).flip(1).cumsum(dim=1).flip(1) + advantages = advantages / lengths.masked_fill(lengths == 0, 1.0) else: advantages = rewards @@ -380,7 +382,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_log_probs, dim=-1, index=shifted_labels.unsqueeze(-1) ).squeeze(-1) - mask = shifted_labels != -100 + # `input_ids` never holds -100; the padding fill lives in `labels`, so the mask has to come from there. The + # gather above keeps indexing `input_ids`, since -100 is not a valid index. + mask = labels[:, prompt_lengths:] != -100 if self.rkl_advantage: reverse_kl_advantage = self._compute_advantage(