Vllm guided decoding (#3663)

* vllm sampling params fix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* do not patch base_trainer

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* seperate vllm fixes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fixup deletion

* Fix indentation

* revert to old style

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

unsloth/__init__.py

@@ -110,6 +110,7 @@ from unsloth_zoo.device_type import (
 from .import_fixes import (
     fix_xformers_performance_issue,
     fix_vllm_aimv2_issue,
+    fix_vllm_guided_decoding_params,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
     patch_trackio,
@@ -118,6 +119,7 @@ from .import_fixes import (
 
 fix_xformers_performance_issue()
 fix_vllm_aimv2_issue()
+fix_vllm_guided_decoding_params()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
 patch_trackio()
@@ -125,6 +127,7 @@ patch_datasets()
 
 del fix_xformers_performance_issue
 del fix_vllm_aimv2_issue
+del fix_vllm_guided_decoding_params
 del ignore_logger_messages
 del patch_ipykernel_hf_xet
 del patch_trackio

unsloth/import_fixes.py

@@ -155,6 +155,22 @@ def fix_vllm_aimv2_issue():
                 print(f"Unsloth: Failed patching vLLM with error = {str(e)}")
 
 
+def fix_vllm_guided_decoding_params():
+    if importlib.util.find_spec("vllm") is None:
+        return
+    # GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM
+    # https://github.com/vllm-project/vllm/pull/22772/files
+    # trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates
+    import vllm
+
+    try:
+        from vllm.sampling_params import GuidedDecodingParams
+    except ImportError:
+        vllm.sampling_params.GuidedDecodingParams = (
+            vllm.sampling_params.StructuredOutputsParams
+        )
+
+
 def ignore_logger_messages():
     # Ignore Environment variable `HF_TOKEN` is set
     try:

unsloth/models/qwen3_moe.py

@@ -59,7 +59,9 @@ def Qwen3MoeSparseMoeBlock_fast_forward(self, X, temp_gate = None, temp_up = Non
         self.gate_proj, X, out = temp_gate
     )  # pretty much the only change from transformers implementation.
 
-    routing_weights = torch_nn_functional_softmax(router_logits, dim = -1, dtype = torch.float32)
+    routing_weights = torch_nn_functional_softmax(
+        router_logits, dim = -1, dtype = torch.float32
+    )
     routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim = -1)
     routing_weights /= routing_weights.sum(dim = -1, keepdim = True)
     # we cast back to the input dtype

unsloth/models/rl.py

@@ -329,6 +329,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     try:
         trainer = eval(f"trl.trainer.{trainer_file}")
     except Exception as error:
+        print(f"Unsloth: Could not import trl.trainer.{trainer_file}: {error}")
         return
 
     # Get SFTTrainer and SFTConfig names
@@ -347,8 +348,14 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         and trainer_file.split("_")[0] in x.lower()
     ]
     if len(name) != 1:
+        print(
+            f"Unsloth: Could not find Trainer class in trl.trainer.{trainer_file}. Found: {name}"
+        )
         return
     if len(config) != 1:
+        print(
+            f"Unsloth: Could not find Config class in trl.trainer.{trainer_file}. Found: {config}"
+        )
         return
 
     # Get SFTTrainer, SFTConfig
@@ -357,16 +364,24 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     try:
         RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}")
     except:
+        print(
+            f"Unsloth: Could not load {RLTrainer_name} from trl.trainer.{trainer_file}"
+        )
         return
     try:
         RLConfig = eval(f"trl.trainer.{trainer_file}.{RLConfig_name}")
     except:
+        print(
+            f"Unsloth: Could not load {RLConfig_name} from trl.trainer.{trainer_file}"
+        )
         return
 
     # Check name
     if RLTrainer.__name__.startswith("Unsloth"):
+        print(f"Unsloth: {RLTrainer.__name__} is already patched.")
         return
     if RLConfig.__name__.startswith("Unsloth"):
+        print(f"Unsloth: {RLConfig.__name__} is already patched.")
         return
 
     # Get old source
@@ -1291,7 +1306,11 @@ def patch_trl_rl_trainers():
     import trl.trainer
 
     all_trainers = dir(trl.trainer)
-    all_trainers = [x for x in all_trainers if x.islower() and x.endswith("_trainer")]
+    all_trainers = [
+        x
+        for x in all_trainers
+        if x.islower() and x.endswith("_trainer") and x != "base_trainer"
+    ]
     for trainer in all_trainers:
         _patch_trl_rl_trainers(trainer)
     return