vision: keep a caller's logits_to_keep on transformers 5

The v5 branch popped logits_to_keep and num_logits_to_keep unconditionally,
so an explicit caller value was discarded before generate() ever saw it.
v5 injects logits_to_keep=1 itself, but that injection is guarded by
`"logits_to_keep" not in model_kwargs`, which makes it a default rather
than an override: a value the caller passed is honored and must not be
dropped.

The cost of dropping it is not a no-op. Measured on a LoRA Qwen2-VL under
transformers 5.14.1, with the model's forward hooked so the validator still
sees the real signature: asking for logits_to_keep=0 (the whole sequence)
reached forward as 1 and returned logits of shape (1, 1, 151936); with the
value preserved it reached forward as 0 and returned (1, 88, 151936).

Deleting the pops outright would be wrong in the other direction. An
explicit num_logits_to_keep raises from _validate_model_kwargs on plain,
PEFT, text and vision models alike, because v5 renamed it away, and
logits_to_keep raises on the 12 of 79 image-text-to-text architectures
whose top-level forward does not take it. So each key is now stripped only
when _unsloth_generate_accepts_kwarg says this model would reject it, which
is the same predicate the validator uses. Under PEFT, self inside the
wrapper is the object the validator later runs against, so the check has no
false negatives there.

Also softened the comment above the branch. Unsloth 2026.7.5 does
pre-inject on a LoRA Qwen2-VL under 5.14.1 and generation succeeds, so
"pre-injecting makes the strict validator raise on PEFT models" overstates
it. Skipping the injection on v5 is still right: it is redundant, and the
arch walk can select a key the top-level model rejects.

tests/test_generate_kwarg_gate.py gains four cases covering a preserved
supported value, a stripped unsupported one, untouched neighbours, and the
absence of the unconditional pop.
This commit is contained in:
Daniel Han 2026-07-26 15:24:13 +00:00
commit bd4ddf3657
2 changed files with 58 additions and 4 deletions

View file

@ -130,6 +130,50 @@ def test_generate_kwarg_gate():
assert got is expected, f"{name}: got {got}, expected {expected}"
# --- v5 logits-to-keep filtering ------------------------------------------
# transformers >= 5 injects logits_to_keep=1 in generate() itself, but the
# injection is guarded by `"logits_to_keep" not in model_kwargs`, so it is a
# DEFAULT. An explicit caller value must survive: popping unconditionally turns
# logits_to_keep=0 (give me the full sequence) into 1 without telling anyone.
# The only values that must be stripped are the ones the strict validator would
# raise on, which is exactly what the gate above predicts.
def _filter_logits_kwargs(model, kwargs):
"""The v5 branch of unsloth_base_fast_generate, as a testable function."""
for key in ("logits_to_keep", "num_logits_to_keep"):
if key in kwargs and not accepts(model, key):
kwargs.pop(key, None)
return kwargs
def test_v5_preserves_a_supported_caller_value():
model = PrepHasKwargs_ForwardHasKey()
# 0 means "all logits"; silently rewriting it to 1 changes the output shape.
assert _filter_logits_kwargs(model, {"logits_to_keep": 0}) == {"logits_to_keep": 0}
assert _filter_logits_kwargs(model, {"logits_to_keep": 5}) == {"logits_to_keep": 5}
def test_v5_strips_a_value_the_model_would_reject():
# num_logits_to_keep was renamed away in v5, so the validator raises on it.
model = PrepHasKwargs_ForwardHasKey()
assert _filter_logits_kwargs(model, {"num_logits_to_keep": 1}) == {}
# A VLM whose top-level forward has no logits_to_keep at all.
assert _filter_logits_kwargs(NoPrepare(), {"logits_to_keep": 1}) == {}
def test_v5_leaves_other_kwargs_alone():
model = PrepHasKwargs_ForwardHasKey()
out = _filter_logits_kwargs(model, {"logits_to_keep": 2, "max_new_tokens": 8})
assert out == {"logits_to_keep": 2, "max_new_tokens": 8}
def test_source_has_no_unconditional_pop():
src = open(VISION).read()
assert 'kwargs.pop("logits_to_keep", None)\n kwargs.pop("num_logits_to_keep", None)' not in src, (
"the v5 branch must not drop caller-supplied logits_to_keep unconditionally"
)
if __name__ == "__main__":
test_generate_kwarg_gate()
for name, _, _, _ in CASES:

View file

@ -399,8 +399,8 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
kwargs.pop("mm_token_type_ids", None)
# VLMs do not allow logits_to_keep. transformers >= 5.0 sets it itself in
# generate() after _validate_model_kwargs, so pre-injecting makes the strict
# validator raise on PEFT models. Skip on v5+ and strip any leaked kwarg.
# generate(), so pre-injecting is redundant there, and the arch walk below
# can pick a key the top-level model rejects. Skip the injection on v5+.
if Version(transformers_version) < Version("5.0.0.dev0"):
global NUM_LOGITS_TO_KEEP
if arch not in NUM_LOGITS_TO_KEEP:
@ -422,8 +422,18 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key):
kwargs[key] = 1
else:
kwargs.pop("logits_to_keep", None)
kwargs.pop("num_logits_to_keep", None)
# v5's own injection (generation/utils.py) is guarded by
# `"logits_to_keep" not in model_kwargs`, so it is a default, not an
# override: an explicit caller value survives and must not be dropped.
# Popping unconditionally silently rewrites logits_to_keep=0 (full
# sequence) into 1. Only strip a key this model would reject, which is
# what the strict validator raises on: num_logits_to_keep everywhere
# (renamed away in v5), and logits_to_keep on the VLMs whose top-level
# forward does not take it.
for _logits_kwarg in ("logits_to_keep", "num_logits_to_keep"):
if _logits_kwarg in kwargs and \
not _unsloth_generate_accepts_kwarg(self, _logits_kwarg):
kwargs.pop(_logits_kwarg, None)
model_eos_token_id = getattr(self.config, "eos_token_id", None)
if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"):