Fix num_logits_to_keep regression on transformers >= 4.52 (#5538)

* Fix num_logits_to_keep on transformers >= 4.51 + compile loss_function

Two follow-ups to the fused-forward work landed in unsloth-zoo PR #665.

1. unsloth_fast_generate (models/llama.py): transformers 4.51 renamed
   num_logits_to_keep to logits_to_keep. Previously we unconditionally
   set kwargs['num_logits_to_keep'] = 1, which transformers 4.57's
   _validate_model_kwargs rejects with:
     ValueError: The following `model_kwargs` are not used by the
     model: ['num_logits_to_keep']
   blocking model.generate() on Llama / Mistral. Now we inspect the
   runtime forward signature and use whichever spelling it accepts;
   if a caller still passes the legacy name we promote it to the new
   spelling instead of stripping it.

2. patch_loss_functions (models/loader.py): the single internal call
   site passed torch_compile=False. UnslothForCausalLMLoss is small
   (label shift + Triton CE), so torch.compile folds the elementwise
   prep into one launch and removes per-step Python overhead. The
   < 2.4 fallback inside patch_loss_functions still routes through
   torch._disable_dynamo so older torches are unaffected.

Verified:
- Llama 3.2 1B + model.generate() no longer raises; emits a sensible
  16-token continuation.
- Gemma3 1B GRPO smoke (max_steps=3) returns bit-identical losses
  0.256 / 0.4393 / 0.2031 vs pre-fix; train_runtime 409s (vs 415s
  pre-fix, within noise).
- unsloth-zoo test_compiler_rewriter_exhaustive + test_fused_forward_install
  pass (96 passed) on this combination.

Related: unslothai/unsloth-zoo PR for the compiler.py single-matmul
backport.

* Revert loader.py loss-compile flip; correct rename-version comment

Drop the patch_loss_functions(torch_compile=True) flip. Tracing the
loss call chain:

  UnslothForCausalLMLoss
    -> unsloth_fixed_cross_entropy
      -> _fast_cross_entropy_loss
         -> Fast_CrossEntropyLoss.apply  (torch.autograd.Function wrapping Triton)

torch.compile treats custom autograd.Function.apply as an opaque op and
breaks the graph at the boundary. The only Python it can actually
compile in the loss function is the label-shift + ignore-fill prep
(three elementwise ops), and the per-call dynamo guard overhead is in
the same order as that prep. Empirical Gemma3 1B GRPO smoke (max_steps=3)
showed no meaningful runtime delta (415s vs 409s, within noise) and
risked dragging the outer compiled training step into recompiles when
the inner guards drift. Keep torch_compile=False; the Triton kernel is
the work, and it is unchanged either way.

Also: the inline comment in unsloth_fast_generate said the kwarg rename
landed in transformers 4.51. The actual decorator (@deprecate_kwarg)
was tagged version="4.50" and present through 4.51.x, then removed in
4.52+. Correct the comment. No behaviour change.
This commit is contained in:
Daniel Han 2026-05-18 02:32:36 -07:00 committed by GitHub
commit 61878c78b8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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
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
elif "num_logits_to_keep" in _fwd_params:
kwargs["num_logits_to_keep"] = 1
# Remove token_type_ids
kwargs.pop("token_type_ids", None)