ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add safe fallbacks for vllm and GRPO edge cases

Browse source
This commit is contained in:
danielhanchen 2026-02-02 02:51:51 +00:00
commit 88c266cefb
7 changed files with 284 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/__init__.py
View file

@ -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

70
unsloth/import_fixes.py
View file

@ -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
)

97
unsloth/models/_utils.py
View file

@ -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]

12
unsloth/models/llama.py
View file

@ -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):

19
unsloth/models/rl.py
View file

@ -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(

80
unsloth/models/rl_replacements.py
View file

@ -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 ")

10
unsloth/models/vision.py
View file

@ -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