Mask padded MiniLLM completion tokens out of the advantage and its length

[behroozazarkhalili]

Fixes #7024. Both defects are in trl/experimental/minillm/minillm_trainer.py and are independent of #6635, which only changes how the discounted sum is computed.

Right-padding leaked into the advantage

compute_loss built the mask it hands to _compute_advantage as input_ids[:, prompt_lengths:] != -100. input_ids never contains -100; the fill is applied to labels, built a few lines earlier from attention_mask. So the mask was all ones, padded slots kept a reward on the pad token, and every earlier position's advantage summed those in. The mask now reads labels[:, prompt_lengths:] != -100. The gather index stays on input_ids: the issue's one-line suggestion of slicing labels for both would feed -100 into torch.gather.

Length normalization depended on batch padding

_compute_advantage replaced masked slots with 1e-4 before the discounted length, so the denominator of a short completion included 1e-4 * gamma^(i-t) for each pad slot after it. With a constant reward of 1.0 and gamma 1.0 a single valid token padded to 512 got 0.9514 instead of 1.0. Only unmasked slots count toward the length now, and clamp(min=1e-4) on the finished length keeps a fully masked row finite, which is what the fill was for.

Tests

• test_length_normalization_ignores_padding: constant reward, gamma 1.0, lengths 1, 128 and 512 padded to 512; the advantage is exactly 1.0 at every valid position and 0 after, and a fully masked row stays finite and zero. Fails on main at lengths 1 and 128.
• test_advantage_mask_excludes_padded_completion_tokens: calls compute_loss on a bare trainer with a stub model and a two-row batch whose first row is half padding, and checks the mask handed to _compute_advantage equals the completion mask. Fails on main, where the mask is all ones.

Both tests are hermetic (no model download) and run in about ten seconds.

Verification

ruff check and ruff format --check at the CI-pinned 0.13.3 and the pinned doc-builder at --max-len 119 pass on both files. On a compute node (job 58038396) the full MiniLLM test file gives 7 passed on the branch, and the two new tests fail on main as described.

---

Note

Medium Risk
Changes experimental training loss/advantage math and can shift MiniLLM optimization dynamics, but the scope is limited to padding handling in minillm_trainer.py.

Overview
Fixes two padding bugs in experimental MiniLLM reverse-KL advantage computation that skewed training when completions were right-padded in a batch.

In compute_loss, the mask passed to _compute_advantage is now derived from labels (where padding is -100) instead of input_ids, so padded completion slots no longer contribute rewards that bleed into earlier token advantages. Log-probability gather still uses input_ids so -100 is never used as an index.

In _compute_advantage, length normalization no longer treats padded positions as tiny contributions via a 1e-4 mask fill; only valid tokens count in the discounted length denominator, with clamp(min=1e-4) on the divisor for fully masked rows. Short completions now get the correct normalized advantage regardless of batch padding length.

Hermetic tests cover constant-reward length normalization across pad lengths and assert compute_loss forwards the completion mask into _compute_advantage.

^{Reviewed by Cursor Bugbot for commit 0cdf67f. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

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.}

[cursor]

Length clamp corrupts late advantages

Medium Severity

The clamp(min=1e-4) on discounted lengths also triggers for valid tokens once gamma^t is small, so later positions get a deflated length-normalized advantage. That hits the long-context RKL path (gamma > 0) at default max_completion_length for common discount values.

 

^{Reviewed by Cursor Bugbot for commit 0cdf67f. Configure here.}

[bot-ci-comment]

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.