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:
parent
47167885b0
commit
61878c78b8
1 changed files with 16 additions and 3 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue