Skip to content

Fix num_logits_to_keep regression on transformers >= 4.52 #5538

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danielhanchen merged 2 commits into from
May 18, 2026
Merged

Fix num_logits_to_keep regression on transformers >= 4.52 #5538

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
482258f
Fix num_logits_to_keep on transformers >= 4.51 + compile loss_function
danielhanchen May 18, 2026
82c2e69
Revert loader.py loss-compile flip; correct rename-version comment
danielhanchen May 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 16 additions & 3 deletions unsloth/models/llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2109,11 +2109,24 @@ def unsloth_fast_generate(

# For newer HF
kwargs["cache_implementation"] = "dynamic"
# For num_logits_to_keep
num_logits_to_keep = kwargs.get("num_logits_to_keep", None)
# transformers 4.50 renamed num_logits_to_keep -> logits_to_keep
# (with @deprecate_kwarg through 4.51.x, removed in 4.52+). Pick the
# spelling the actual runtime forward accepts so generation
# _validate_model_kwargs does not reject the legacy name.
num_logits_to_keep = kwargs.pop("num_logits_to_keep", None)
logits_to_keep = kwargs.get("logits_to_keep", None)
if num_logits_to_keep is not None and logits_to_keep is None:
kwargs["logits_to_keep"] = num_logits_to_keep
Comment on lines +2118 to +2119

@chatgpt-codex-connector chatgpt-codex-connector Bot May 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve legacy logits kwarg on old transformers

When a caller supplies num_logits_to_keep on transformers versions whose model forward still only accepts that legacy spelling, this code always pops it and forwards logits_to_keep instead. Generation then hits HF's model-kwarg validation with an unexpected logits_to_keep, so existing callers on pre-rename transformers regress; the translation needs to be based on the inspected forward signature before replacing the kwarg.

Useful? React with 👍 / 👎.

logits_to_keep = num_logits_to_keep
if num_logits_to_keep is None and logits_to_keep is None:
kwargs["num_logits_to_keep"] = 1
try:
_fwd_params = inspect.signature(self.forward).parameters
except (TypeError, ValueError):
_fwd_params = {}
if "logits_to_keep" in _fwd_params:
kwargs["logits_to_keep"] = 1
Comment on lines +2126 to +2127

@chatgpt-codex-connector chatgpt-codex-connector Bot May 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve Mistral's logits-slicing kwarg

When generate() runs on FastMistralModel with no explicit logits kwarg, this branch now chooses logits_to_keep because Mistral's patched forward signature contains both names. However MistralForCausalLM_fast_forward only uses num_logits_to_keep in its normal logits path (unsloth/models/mistral.py:316) and only merges the two names for the hidden-states env path, so Mistral generation falls back to computing full prompt logits instead of slicing to the last token. For long prompts this reintroduces the large prefill logits allocation that the default =1 was avoiding and can cause major slowdown/OOMs.

Useful? React with 👍 / 👎.

elif "num_logits_to_keep" in _fwd_params:
kwargs["num_logits_to_keep"] = 1
Comment on lines +2116 to +2129

@gemini-code-assist gemini-code-assist Bot May 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current logic for migrating num_logits_to_keep to logits_to_keep breaks backward compatibility for older transformers versions when the parameter is explicitly passed by the user.

Specifically, if a user provides num_logits_to_keep on a version of transformers that does not yet support logits_to_keep, the code at lines 2117-2119 will unconditionally rename the key in kwargs. This will cause transformers' _validate_model_kwargs to raise a ValueError because logits_to_keep is not in the model's forward signature.

Furthermore, logits_to_keep is not popped from kwargs at line 2116, which could lead to similar validation errors if the model only accepts the legacy name.

A unified approach that normalizes both inputs and then uses signature inspection to decide which key to use is more robust.

    provided_num = kwargs.pop("num_logits_to_keep", None)
    provided_logits = kwargs.pop("logits_to_keep", None)
    val = provided_logits if provided_logits is not None else provided_num

    try:
        _fwd_params = inspect.signature(self.forward).parameters
    except (TypeError, ValueError):
        _fwd_params = {}

    if "logits_to_keep" in _fwd_params:
        kwargs["logits_to_keep"] = val if val is not None else 1
    elif "num_logits_to_keep" in _fwd_params:
        kwargs["num_logits_to_keep"] = val if val is not None else 1


# Remove token_type_ids
kwargs.pop("token_type_ids", None)
Expand Down
Loading