diff --git a/tests/test_fast_generate_slow_guard.py b/tests/test_fast_generate_slow_guard.py new file mode 100644 index 000000000..6bfc561e5 --- /dev/null +++ b/tests/test_fast_generate_slow_guard.py @@ -0,0 +1,95 @@ +"""GPU-free test for the fast_generate slow-mode guard in _utils.py. + +When fast_inference=False, model.fast_generate falls back to HuggingFace generate, so vLLM-only +inputs must be rejected with a clear message instead of leaking into transformers.generate. Covers +a string prompt, a vLLM {"prompt":..., "multi_modal_data":...} dict, SamplingParams passed both +positionally and as a kwarg, and a normal tokenized call passing through. +""" + +import ast, functools, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UTILS = os.path.join(HERE, "unsloth", "models", "_utils.py") + + +def _load_factory(): + src = open(UTILS).read() + for node in ast.parse(src).body: + if isinstance(node, ast.FunctionDef) and node.name == "make_fast_generate_wrapper": + ns = {"functools": functools} + exec(ast.get_source_segment(src, node), ns) + return ns["make_fast_generate_wrapper"] + raise AssertionError("make_fast_generate_wrapper not found in _utils.py") + + +make_fast_generate_wrapper = _load_factory() + + +class _SamplingParams: + pass + + +_SamplingParams.__name__ = "SamplingParams" # match by class name, no vllm import needed + + +def _wrapper(): + state = {} + + def original_generate(*a, **k): + state["hit"] = True + return "ok" + + return make_fast_generate_wrapper(original_generate), state + + +def _rejects(fn, needle): + try: + fn() + except ValueError as e: + assert needle in str(e), str(e) + return True + raise AssertionError("expected ValueError") + + +def test_fast_generate_slow_guard(): + w, _ = _wrapper() + # reject every vLLM-only shape + assert _rejects(lambda: w("hello"), "fast_inference=True") + assert _rejects( + lambda: w({"prompt": "hi", "multi_modal_data": {"image": None}}), "fast_inference=True" + ) + assert _rejects(lambda: w(["a", "b"]), "fast_inference=True") + assert _rejects(lambda: w([{"prompt": "hi"}]), "fast_inference=True") # list of prompt dicts + assert _rejects( + lambda: w({"prompt_token_ids": [1, 2, 3]}), "fast_inference=True" + ) # vLLM TokensPrompt + assert _rejects(lambda: w(prompts = "hello"), "fast_inference=True") # vLLM `prompts` kwarg + assert _rejects( + lambda: w(prompts = [{"prompt": "hi"}]), "fast_inference=True" + ) # vLLM `prompts` kwarg list + assert _rejects( + lambda: w(prompt_token_ids = [1, 2, 3]), "fast_inference=True" + ) # vLLM legacy tokenized kwarg + assert _rejects( + lambda: w(prompts = [1, 2, 3]), "fast_inference=True" + ) # token-id list via vLLM-only `prompts` kwarg + assert _rejects( + lambda: w(prompts = None), "fast_inference=True" + ) # vLLM-only kwarg present even if None + assert _rejects(lambda: w({"prompt": "hi"}, _SamplingParams()), "sampling_params") + assert _rejects( + lambda: w({"prompt": "hi"}, [_SamplingParams()]), "sampling_params" + ) # list of SamplingParams + assert _rejects(lambda: w(sampling_params = object()), "sampling_params") + + # pass normal tokenized calls with no false positives + w, state = _wrapper() + assert w(input_ids = "TOKENS", max_new_tokens = 8) == "ok" and state.get("hit") + assert w([1, 2, 3], max_new_tokens = 8) == "ok" # positional token ids + assert w([], max_new_tokens = 8) == "ok" # empty positional + print("13 reject + 3 pass fast_generate slow-mode guard cases passed") + + +if __name__ == "__main__": + test_fast_generate_slow_guard() + print("OK: fast_generate rejects vLLM-style inputs when fast_inference=False") diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 599a5c026..047783c35 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3602,8 +3602,27 @@ def make_fast_generate_wrapper(original_generate): @functools.wraps(original_generate) def _fast_generate_wrapper(*args, **kwargs): - # Check for vLLM-specific arguments - if "sampling_params" in kwargs: + def _has_sampling_params(a): + # SamplingParams passed directly or inside a positional list/tuple + return type(a).__name__ == "SamplingParams" or ( + isinstance(a, (list, tuple)) + and any(type(i).__name__ == "SamplingParams" for i in a) + ) + + def _is_vllm_prompt(a): + # str prompt, a vLLM prompt dict (prompt / prompt_token_ids / prompt_embeds / + # multi_modal_data), or a list/tuple of those + head = a[0] if isinstance(a, (list, tuple)) and len(a) > 0 else a + return isinstance(head, str) or ( + isinstance(head, dict) + and any( + k in head + for k in ("prompt", "prompt_token_ids", "prompt_embeds", "multi_modal_data") + ) + ) + + # vLLM-only; also catch SamplingParams passed positionally (fast_generate(prompt, params)) + if "sampling_params" in kwargs or any(_has_sampling_params(a) for a in args): raise ValueError( "Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). " "Since `fast_inference=False`, use HuggingFace generate arguments instead:\n" @@ -3616,33 +3635,26 @@ def make_fast_generate_wrapper(original_generate): "Since `fast_inference=False`, LoRA weights are already merged into the model." ) - # Check if first positional argument is a string or list of strings - if len(args) > 0: - first_arg = args[0] - is_string_input = False - - if isinstance(first_arg, str): - is_string_input = True - elif isinstance(first_arg, (list, tuple)) and len(first_arg) > 0: - if isinstance(first_arg[0], str): - is_string_input = True - - if is_string_input: - raise ValueError( - "Unsloth: Passing text strings to `fast_generate` is only supported " - "when `fast_inference=True` (vLLM). Since `fast_inference=False`, you must " - "tokenize the input first:\n\n" - " messages = tokenizer.apply_chat_template(\n" - ' [{"role": "user", "content": "Your prompt here"}],\n' - " tokenize=True, add_generation_prompt=True,\n" - ' return_tensors="pt", return_dict=True\n' - " )\n" - " output = model.fast_generate(\n" - " **messages.to('cuda'),\n" - " max_new_tokens=64,\n" - " temperature=1.0,\n" - " )" - ) + # A vLLM-style prompt (string, {"prompt":..., "multi_modal_data":...} dict, or a list/tuple + # of either) only works under vLLM; tokenize first when fast_inference=False. A positional + # arg may be HF token ids, so check it conservatively with _is_vllm_prompt. The `prompts` / + # `prompt_token_ids` / `prompt_embeds` keywords are vLLM-only names that HuggingFace generate + # does not accept, so any of them being present is a vLLM-style call (even a bare token list, + # or an explicit None from a defaulted kwargs dict), hence membership rather than a value check. + vllm_prompt_kwarg = any( + k in kwargs for k in ("prompts", "prompt_token_ids", "prompt_embeds") + ) + if (len(args) > 0 and _is_vllm_prompt(args[0])) or vllm_prompt_kwarg: + raise ValueError( + "Unsloth: Passing vLLM-style prompts to `fast_generate` is only supported when " + "`fast_inference=True` (vLLM). Since `fast_inference=False`, tokenize first:\n\n" + " inputs = tokenizer.apply_chat_template(\n" + ' [{"role": "user", "content": "Your prompt here"}],\n' + " tokenize=True, add_generation_prompt=True,\n" + ' return_tensors="pt", return_dict=True,\n' + " )\n" + " output = model.fast_generate(**inputs.to('cuda'), max_new_tokens=64, temperature=1.0)" + ) # Call original generate return original_generate(*args, **kwargs)