Repository navigation
Mask padded MiniLLM completion tokens out of the advantage and its length - #7044
behroozazarkhalili wants to merge 5 commits into
Conversation
…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
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 0cdf67f. Configure here.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
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.
|
Closing this one myself. It came from reading the code, not from a failure in a real training run, and per the discussion on #6862 that isn't the bar for this repo. If someone hits this in practice, the branch is still here and I'm happy to reopen it with a reproduction. |
1 participant
MiniLLM included padded tokens in advantages and let padding change the length normalization denominator. I built the advantage mask from labels and counted only unmasked positions in the discounted length, with a guard for empty positions. Both new regression tests fail on main, and the full MiniLLM test file passes on this branch.
Fixes #7024.
Note
Medium Risk
Changes experimental MiniLLM training math (advantage masking and normalization), which directly affects loss gradients but is a targeted correctness fix with new regression tests.
Overview
Fixes incorrect reverse-KL advantages in MiniLLM when batches use right-padded completions.
compute_lossnow builds the completion mask fromlabels(where padding is-100) instead ofshifted_labelsfrominput_ids, so padded completion slots no longer contribute rewards or bleed into earlier positions’ advantages. Log-prob gathers still useinput_idsindices._compute_advantagelength normalization no longer treats masked positions as1e-4in the discounted-length denominator. Only unmasked tokens count, so short completions padded to a fixed width get the correct normalized advantage (e.g. 1.0 for constant reward withgamma=1.0). Zero-length tails divide by 1 instead of using a floor that would distort late valid tokens whengamma < 1.Adds hermetic tests for padding-invariant length normalization and for the mask passed into
_compute_advantagefromcompute_loss.Reviewed by Cursor Bugbot for commit 9ccdf8f. Bugbot is set up for automated code reviews on this repo. Configure here.