diff --git a/tests/test_generate_kwarg_gate.py b/tests/test_generate_kwarg_gate.py index 6d1379d3a9..242fcebfbf 100644 --- a/tests/test_generate_kwarg_gate.py +++ b/tests/test_generate_kwarg_gate.py @@ -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: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index f20f32cc3e..effe633302 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -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__"):