From 9db2e12252b9a4b32e87abc9041f13ca00343fc2 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Thu, 7 Nov 2024 01:52:08 +0530 Subject: [PATCH] Throw error when inferencing longer than max_popsition_embeddings (#1236) * Throw error when inferencing longer than max_popsition_embeddings without rope scaling * Update llama.py --------- Co-authored-by: Daniel Han --- unsloth/models/llama.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 3c4d8f3b38..7f07bea4c5 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1376,6 +1376,15 @@ def _wrap_fast_inference(generate, device_type, dtype, model): @torch.inference_mode def _fast_generate(*args, **kwargs): + if hasattr(model, "config") and hasattr(model.config, "max_position_embeddings"): + if "input_ids" in kwargs and kwargs["input_ids"] is not None and "max_new_tokens" in kwargs: + if kwargs["input_ids"].shape[-1] + kwargs["max_new_tokens"] > model.config.max_position_embeddings: + raise ValueError( + f'Unsloth: input length {kwargs["input_ids"].shape[-1]} + max_new_tokens {kwargs["max_new_tokens"]} exceeds the maximum sequence length of {model.config.max_position_embeddings}!\n'\ + 'You will need to do long context extension by increasing the `max_seq_length` in `FastLanguageModel.from_pretrained`.' + ) + pass + # Set a flag for generation! internal_model = model while hasattr(internal_model, "model"):