From 88c266cefb46f2dce91d949db5dc2216db0054b9 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 2 Feb 2026 02:51:51 +0000 Subject: [PATCH] Add safe fallbacks for vllm and GRPO edge cases --- unsloth/__init__.py | 9 +++ unsloth/import_fixes.py | 70 ++++++++++++++++++++-- unsloth/models/_utils.py | 97 +++++++++++++++++++++++++++++++ unsloth/models/llama.py | 12 +++- unsloth/models/rl.py | 19 +++--- unsloth/models/rl_replacements.py | 80 +++++++++++++++++++++++++ unsloth/models/vision.py | 10 +++- 7 files changed, 284 insertions(+), 13 deletions(-) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 0b819a546e..668a930852 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -124,6 +124,9 @@ from unsloth_zoo.device_type import ( # Fix other issues from .import_fixes import ( fix_xformers_performance_issue, + fix_deepseek_v2_moe_alias, + fix_qwen2_vl_max_pixels_none, + fix_qwen2_5_vl_tie_word_embeddings, fix_vllm_aimv2_issue, fix_vllm_guided_decoding_params, fix_vllm_pdl_blackwell, @@ -139,6 +142,9 @@ from .import_fixes import ( ) fix_xformers_performance_issue() +fix_deepseek_v2_moe_alias() +fix_qwen2_vl_max_pixels_none() +fix_qwen2_5_vl_tie_word_embeddings() fix_vllm_aimv2_issue() fix_vllm_guided_decoding_params() fix_vllm_pdl_blackwell() @@ -153,6 +159,9 @@ fix_executorch() patch_vllm_for_notebooks() del fix_xformers_performance_issue +del fix_deepseek_v2_moe_alias +del fix_qwen2_vl_max_pixels_none +del fix_qwen2_5_vl_tie_word_embeddings del fix_vllm_aimv2_issue del fix_vllm_guided_decoding_params del fix_vllm_pdl_blackwell diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 89fd152857..27f97a7e92 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -204,6 +204,60 @@ def fix_xformers_performance_issue(): logger.info(f"Unsloth: Failed patching Xformers with error = {str(e)}") +def fix_deepseek_v2_moe_alias(): + try: + from transformers.models.deepseek_v2 import modeling_deepseek_v2 as deepseek_v2 + except Exception: + return + + if hasattr(deepseek_v2, "DeepseekV2MoE"): + return + if not hasattr(deepseek_v2, "DeepseekV2Moe"): + return + + deepseek_v2.DeepseekV2MoE = deepseek_v2.DeepseekV2Moe + try: + if hasattr(deepseek_v2, "__all__") and "DeepseekV2MoE" not in deepseek_v2.__all__: + deepseek_v2.__all__.append("DeepseekV2MoE") + except Exception: + pass + logger.info("Unsloth: Added DeepseekV2MoE alias for DeepseekV2Moe") + + +def fix_qwen2_vl_max_pixels_none(): + try: + from transformers.models.qwen2_vl import image_processing_qwen2_vl as qwen2_vl + except Exception: + return + + if getattr(qwen2_vl.smart_resize, "_unsloth_max_pixels_patch", False): + return + + original = qwen2_vl.smart_resize + default_max_pixels = 14 * 14 * 4 * 1280 + + def smart_resize(height, width, factor=28, min_pixels=56 * 56, max_pixels=default_max_pixels): + if max_pixels is None: + max_pixels = default_max_pixels + return original(height, width, factor=factor, min_pixels=min_pixels, max_pixels=max_pixels) + + smart_resize._unsloth_max_pixels_patch = True + qwen2_vl.smart_resize = smart_resize + + +def fix_qwen2_5_vl_tie_word_embeddings(): + try: + from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLTextConfig + except Exception: + return + + if hasattr(Qwen2_5_VLTextConfig, "tie_word_embeddings"): + return + + Qwen2_5_VLTextConfig.tie_word_embeddings = True + logger.info("Unsloth: Added tie_word_embeddings to Qwen2_5_VLTextConfig") + + def patch_vllm_for_notebooks(): import sys @@ -332,18 +386,26 @@ def fix_vllm_guided_decoding_params(): # trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates try: import vllm - except ImportError as e: + except Exception as e: _maybe_raise_vllm_transformers_mismatch(e) - raise + logger.warning( + "Unsloth: vLLM import failed, skipping GuidedDecodingParams patch. " + f"Error: {e}" + ) + return try: from vllm.sampling_params import GuidedDecodingParams - except ImportError as e: + except Exception as e: _maybe_raise_vllm_transformers_mismatch(e) if not hasattr(vllm, "sampling_params") or not hasattr( vllm.sampling_params, "StructuredOutputsParams" ): - raise + logger.warning( + "Unsloth: vLLM sampling_params missing StructuredOutputsParams; " + "skipping GuidedDecodingParams patch." + ) + return vllm.sampling_params.GuidedDecodingParams = ( vllm.sampling_params.StructuredOutputsParams ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 76952b00a5..5985db04a4 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -74,6 +74,7 @@ __all__ = [ "_get_inference_mode_context_manager", "hf_login", "make_fast_generate_wrapper", + "make_vllm_fast_generate_wrapper", ] import torch @@ -2482,3 +2483,99 @@ def make_fast_generate_wrapper(original_generate): return original_generate(*args, **kwargs) return _fast_generate_wrapper + + +def make_vllm_fast_generate_wrapper(model, vllm_generate): + """ + Wraps vLLM generate to optionally fall back to HF generate on failure. + Keeps vLLM for training; fallback only when model.training is False. + """ + + @functools.wraps(vllm_generate) + def _vllm_fast_generate_wrapper(*args, **kwargs): + if (not model.training) and ( + getattr(model, "_unsloth_disable_vllm_inference", False) + or os.environ.get("UNSLOTH_VLLM_DISABLE_INFERENCE", "0") == "1" + ): + return _vllm_fallback_to_hf_generate(model, args, kwargs) + if getattr(model, "_unsloth_vllm_failed", False): + return _vllm_fallback_to_hf_generate(model, args, kwargs) + try: + return vllm_generate(*args, **kwargs) + except Exception: + if model.training: + raise + if os.environ.get("UNSLOTH_VLLM_FALLBACK", "1") != "1": + raise + model._unsloth_vllm_failed = True + try: + model.fast_generate = make_fast_generate_wrapper(model.generate) + model.fast_generate_batches = None + except Exception: + pass + return _vllm_fallback_to_hf_generate(model, args, kwargs) + + return _vllm_fast_generate_wrapper + + +def _vllm_fallback_to_hf_generate(model, args, kwargs): + tokenizer = getattr(model, "_saved_temp_tokenizer", None) + hf_kwargs = dict(kwargs) + sampling_params = hf_kwargs.pop("sampling_params", None) + hf_kwargs.pop("lora_request", None) + + if sampling_params is not None: + for key in ("temperature", "top_p", "top_k"): + value = getattr(sampling_params, key, None) + if value is not None and key not in hf_kwargs: + hf_kwargs[key] = value + max_tokens = getattr(sampling_params, "max_tokens", None) + if max_tokens is None: + max_tokens = getattr(sampling_params, "max_new_tokens", None) + if max_tokens is not None and "max_new_tokens" not in hf_kwargs: + hf_kwargs["max_new_tokens"] = max_tokens + + if len(args) > 0: + first_arg = args[0] + if isinstance(first_arg, str) or ( + isinstance(first_arg, (list, tuple)) + and len(first_arg) > 0 + and isinstance(first_arg[0], str) + ): + if tokenizer is None: + raise RuntimeError( + "Unsloth: vLLM fast_generate failed and no tokenizer was cached for HF fallback." + ) + inputs = tokenizer( + first_arg, + return_tensors = "pt", + padding = True, + ) + device = getattr(model, "device", None) or "cuda" + inputs = inputs.to(device) + args = () + hf_kwargs = {**inputs, **hf_kwargs} + + outputs = model.generate(*args, **hf_kwargs) + + if tokenizer is None: + return outputs + + try: + texts = tokenizer.batch_decode(outputs, skip_special_tokens = True) + except Exception: + return outputs + + class _FallbackOutput: + __slots__ = ("text",) + + def __init__(self, text): + self.text = text + + class _FallbackRequestOutput: + __slots__ = ("outputs",) + + def __init__(self, text): + self.outputs = [_FallbackOutput(text)] + + return [_FallbackRequestOutput(text) for text in texts] diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 39f2ba1460..bb31b95a32 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2417,10 +2417,14 @@ class FastLlamaModel: quant_state_dict, model_config, dtype, bnb_config ) model.vllm_engine = llm - model.fast_generate = model.vllm_engine.generate + model.fast_generate = make_vllm_fast_generate_wrapper( + model, model.vllm_engine.generate + ) model.fast_generate_batches = functools.partial( generate_batches, model.vllm_engine ) + if isinstance(model_name, str) and "fp8" in model_name.lower(): + model._unsloth_disable_vllm_inference = True raise_handler.remove() # Return old flag os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer @@ -2438,6 +2442,12 @@ class FastLlamaModel: model, tokenizer = patch_tokenizer(model, tokenizer) model, tokenizer = model_patcher.post_patch(model, tokenizer) + if fast_inference and hasattr(model, "vllm_engine"): + try: + model.vllm_engine._unsloth_hf_model = model + model.vllm_engine._unsloth_tokenizer = tokenizer + except Exception: + pass # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 803153e608..daded6fe00 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1131,13 +1131,18 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): RLTrainer_source = re.sub(r"[\n]{3,}", "\n", RLTrainer_source) # Create new function - created_module = create_new_function( - f"Unsloth{RLTrainer_name}", - RLTrainer_source, - f"trl.trainer.{trainer_file}", - imports, - overwrite = False, - ) + try: + created_module = create_new_function( + f"Unsloth{RLTrainer_name}", + RLTrainer_source, + f"trl.trainer.{trainer_file}", + imports, + overwrite = False, + ) + except Exception as exc: + if os.environ.get("UNSLOTH_LOGGING_ENABLED", "0") == "1": + print(f"Unsloth: Failed to compile {RLTrainer_name} ({exc}), falling back to original trainer.") + return # Patch Trainer exec( diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index ab0c37a608..62985746f4 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -434,6 +434,86 @@ def grpo_trainer__generate_and_score_completions(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_and_score_completions) +# Normalize non-string completions for reward functions +def grpo_trainer__calculate_rewards(function_name, function): + if function_name != "_calculate_rewards": + return function + + function = function.replace( + " rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device)\n", + " rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device)\n" + " def _unsloth_completion_to_text(completion):\n" + " if isinstance(completion, str):\n" + " return completion\n" + " if isinstance(completion, dict):\n" + " if completion.get('text') is not None:\n" + " return completion['text']\n" + " if completion.get('content') is not None:\n" + " return completion['content']\n" + " if completion.get('message') is not None:\n" + " return _unsloth_completion_to_text(completion['message'])\n" + " return str(completion)\n" + " if isinstance(completion, list):\n" + " return ''.join(_unsloth_completion_to_text(item) for item in completion)\n" + " return str(completion)\n" + " completion_texts = [_unsloth_completion_to_text(c) for c in completions]\n" + " completions_are_text = all(isinstance(c, str) for c in completions)\n" + ) + + function = function.replace( + " reward_kwargs[\"trainer_state\"] = self.state\n", + " reward_kwargs[\"trainer_state\"] = self.state\n" + " reward_kwargs[\"completion_texts\"] = completion_texts\n" + " reward_kwargs[\"completion_raw\"] = completions\n", + ) + + function = function.replace( + " texts = [p + c for p, c in zip(prompts, completions, strict=True)]\n", + " texts = [p + c for p, c in zip(prompts, (completions if completions_are_text else completion_texts), strict=True)]\n", + ) + + # Add a robust try/except wrapper for reward funcs expecting dict completions. + if "string indices must be integers" not in function and "output_reward_func = reward_func(" in function: + base_try = ( + " # UNSLOTH_REWARD_FUNC_TRY\n" + " try:\n" + " output_reward_func = reward_func(\n" + " prompts=prompts, completions=(completions if completions_are_text else completion_texts), completion_ids=completion_ids_list, **reward_kwargs\n" + " )\n" + " except TypeError as e:\n" + " if \"string indices must be integers\" in str(e):\n" + " def _wrap_completion_list(_comps):\n" + " if isinstance(_comps, list):\n" + " return [c if isinstance(c, dict) else {\"content\": c} for c in _comps]\n" + " return [{\"content\": _comps}]\n" + " _wrapped = [_wrap_completion_list(c) for c in completions]\n" + " output_reward_func = reward_func(\n" + " prompts=prompts, completions=_wrapped, completion_ids=completion_ids_list, **reward_kwargs\n" + " )\n" + " else:\n" + " raise\n" + ) + variant_a = ( + " output_reward_func = reward_func(\n" + " prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs\n" + " )\n" + ) + variant_b = ( + " output_reward_func = reward_func(\n" + " prompts=prompts, completions=(completions if completions_are_text else completion_texts), completion_ids=completion_ids_list, **reward_kwargs\n" + " )\n" + ) + if variant_a in function: + function = function.replace(variant_a, base_try) + elif variant_b in function: + function = function.replace(variant_b, base_try) + + return function + + +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__calculate_rewards) + + # Fix {"reasoning_effort" : "high"} not applied def grpo_trainer_fix_maybe_apply_chat_template(function_name, function): spaces = function.find("def ") diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 6835f2e986..9e1de46b51 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -767,7 +767,9 @@ class FastBaseModel: is_vision_model = is_vlm, ) model.vllm_engine = llm - model.fast_generate = model.vllm_engine.generate + model.fast_generate = make_vllm_fast_generate_wrapper( + model, model.vllm_engine.generate + ) model.fast_generate_batches = functools.partial( generate_batches, model.vllm_engine ) @@ -851,6 +853,12 @@ class FastBaseModel: correct_dtype = correct_dtype, ) model, tokenizer = patch_tokenizer(model, tokenizer) + if fast_inference and hasattr(model, "vllm_engine"): + try: + model.vllm_engine._unsloth_hf_model = model + model.vllm_engine._unsloth_tokenizer = tokenizer + except Exception: + pass model = post_patch_loss_function(model) # Log Unsloth version for future fastpaths for inference