From 0cdf67ff0e94425b214ce20a5c5ffb7fc283eeb0 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Thu, 3 Sep 2026 21:22:36 -0700 Subject: [PATCH 1/2] fix(minillm): mask padded completion tokens out of the advantage and out of its length `compute_loss` built the advantage mask from `input_ids`, which never holds -100, so every padded slot kept a reward of `teacher_logp - student_logp` on the pad token and every earlier position summed it in. The mask now comes from `labels`, where the -100 fill lives; the gather keeps indexing `input_ids`, since -100 is not a valid index. Under `length_normalization=True`, `_compute_advantage` replaced masked slots with 1e-4 before building the discounted length, so a short completion's denominator grew with the longest completion in its batch (a single valid token padded to 512 got 0.9514 instead of 1.0 for a constant reward). Only unmasked slots count now, and a clamp on the finished length keeps a fully masked row finite. Tests: constant reward gives exactly 1.0 at every valid position for lengths 1, 128 and 512 padded to 512; a fully masked row stays finite and zero; the mask handed to `_compute_advantage` by `compute_loss` equals the completion mask. Fixes #7024 --- tests/experimental/test_minillm_trainer.py | 57 +++++++++++++++++++++ trl/experimental/minillm/minillm_trainer.py | 12 +++-- 2 files changed, 64 insertions(+), 5 deletions(-) diff --git a/tests/experimental/test_minillm_trainer.py b/tests/experimental/test_minillm_trainer.py index 59c0be9f306..db5d39325e1 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,58 @@ 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).""" + stub = types.SimpleNamespace(gamma=1.0, 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) + 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..cf3c783eb61 100644 --- a/trl/experimental/minillm/minillm_trainer.py +++ b/trl/experimental/minillm/minillm_trainer.py @@ -336,10 +336,10 @@ 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. The clamp keeps a fully masked row finite. + lengths = (mask * gamma_pow).flip(1).cumsum(dim=1).flip(1) + advantages = advantages / lengths.clamp(min=1e-4) else: advantages = rewards @@ -380,7 +380,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( From 7a573e607d01ad4c43cff1311a78220e81f392f0 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Thu, 3 Sep 2026 22:23:25 -0700 Subject: [PATCH 2/2] fix(minillm): guard only the zero-length tail in length normalization The previous commit floored the discounted length at 1e-4 to keep a fully masked row finite. With gamma below 1 the discounted length of a late valid position is far smaller than that floor (0.9**300 is 1.9e-14), so the floor deflated those positions: with a constant reward of 1.0 and gamma 0.9 position 300 of an unpadded 512-token row got 1.9e-9 instead of 1.0. A length is zero only where every later position is masked, and there the numerator is zero as well, so the guard now replaces zero lengths with one and leaves every valid position alone. The padding test now runs at gamma 1.0 and 0.9 and asserts the advantage equals the reward at every valid position of rows of length 1, 128 and 512; the floored version fails it at 19 of 128 positions. --- tests/experimental/test_minillm_trainer.py | 30 ++++++++++++--------- trl/experimental/minillm/minillm_trainer.py | 6 +++-- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/tests/experimental/test_minillm_trainer.py b/tests/experimental/test_minillm_trainer.py index db5d39325e1..619e705dd46 100644 --- a/tests/experimental/test_minillm_trainer.py +++ b/tests/experimental/test_minillm_trainer.py @@ -88,18 +88,24 @@ 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).""" - stub = types.SimpleNamespace(gamma=1.0, 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) - 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)) + 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: diff --git a/trl/experimental/minillm/minillm_trainer.py b/trl/experimental/minillm/minillm_trainer.py index cf3c783eb61..e37a8af5882 100644 --- a/trl/experimental/minillm/minillm_trainer.py +++ b/trl/experimental/minillm/minillm_trainer.py @@ -337,9 +337,11 @@ def _compute_advantage( if self.length_normalization: # Only unmasked positions count toward the discounted length, so the advantage of a completion does - # not depend on how much padding the batch carries. The clamp keeps a fully masked row finite. + # 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.clamp(min=1e-4) + advantages = advantages / lengths.masked_fill(lengths == 0, 1.0) else: advantages = rewards