Fix GRPO resume on newer TRL/vLLM stack

unsloth/kernels/fast_lora.py

@@ -25,6 +25,10 @@ from .utils import (
 )
 
 
+def _match_backward_weight_dtype(weight, dtype):
+    return weight if weight.dtype == dtype else weight.to(dtype)
+
+
 class LoRA_MLP(torch.autograd.Function):
     """
     ### LoRA weights
@@ -191,12 +195,14 @@ class LoRA_MLP(torch.autograd.Function):
         # dX  = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS)
         # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS)
         upW = fast_dequantize(upW.t(), upW_quant)
+        upW = _match_backward_weight_dtype(upW, df.dtype)
         dX = torch.matmul(df, upW.t(), out = X if ctx.inplace else None)
         del upW
         # dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t())
         dX.addmm_(df @ upB.t(), upA.t(), alpha = upS)
 
         gateW = fast_dequantize(gateW.t(), gateW_quant)
+        gateW = _match_backward_weight_dtype(gateW, dX.dtype)
         # dX += de @ gateW.t()
         dX.addmm_(de, gateW.t())
         del gateW
@@ -487,6 +493,7 @@ class LoRA_QKV(torch.autograd.Function):
         # Combine derivatives to find dX
         # dQ
         QW = fast_dequantize(QW.t(), QW_quant)
+        QW = _match_backward_weight_dtype(QW, dQ.dtype)
         dX = torch.matmul(dQ, QW.t(), out = X if ctx.inplace else None)
         del QW
         # dX += (dQ @ QB.to(dtype).t() @ (QS * QA.to(dtype).t()))
@@ -494,6 +501,7 @@ class LoRA_QKV(torch.autograd.Function):
 
         # dK
         KW = fast_dequantize(KW.t(), KW_quant)
+        KW = _match_backward_weight_dtype(KW, dX.dtype)
         # dX += dK @ KW.t()
         dX.addmm_(dK, KW.t())
         del KW
@@ -502,6 +510,7 @@ class LoRA_QKV(torch.autograd.Function):
 
         # dV
         VW = fast_dequantize(VW.t(), VW_quant)
+        VW = _match_backward_weight_dtype(VW, dX.dtype)
         # dX += dV @ VW.t()
         dX.addmm_(dV, VW.t())
         del VW
@@ -629,6 +638,7 @@ class LoRA_W(torch.autograd.Function):
 
         # Get derivative for dX
         W = fast_dequantize(W.t(), W_quant)
+        W = _match_backward_weight_dtype(W, dY.dtype)
         dX = dY @ W.t()
         del W
         # dX += dY @ B.to(dtype).t() @ (S * A.to(dtype).t())

unsloth/models/llama.py

@@ -2479,6 +2479,8 @@ class FastLlamaModel:
             llm = load_vllm(**load_vllm_kwargs)
 
             # Convert to HF format
+            if getattr(model_config, "model_name", None) is None:
+                model_config.model_name = model_name
             _, quant_state_dict = get_vllm_state_dict(
                 llm,
                 config = model_config,

unsloth/models/rl.py

@@ -156,6 +156,62 @@ def _patch_resume_from_checkpoint_memory(trainer_class):
     trainer_class.train = _unsloth_train_with_resume_guard
 
 
+
+def _maybe_prepare_vllm_for_resume(trainer):
+    if not torch.cuda.is_available():
+        return
+
+    llm = getattr(trainer, "llm", None)
+    if llm is None:
+        llm = getattr(getattr(trainer, "model", None), "vllm_engine", None)
+
+    slept = False
+    sleep_fn = getattr(llm, "sleep", None)
+    if callable(sleep_fn):
+        try:
+            sleep_mode = int(os.environ.get("VLLM_SLEEP_MODE", "1"))
+            sleep_fn(sleep_mode)
+            slept = True
+        except TypeError:
+            try:
+                sleep_fn()
+                slept = True
+            except Exception:
+                pass
+        except Exception:
+            pass
+
+    if slept:
+        trainer._unsloth_resume_wake_vllm = True
+
+    import gc
+    for _ in range(3):
+        gc.collect()
+        torch.cuda.empty_cache()
+pass
+
+
+def _patch_resume_from_checkpoint_memory(trainer_class):
+    original_train = getattr(trainer_class, "train", None)
+    if original_train is None:
+        return
+    if getattr(original_train, "_unsloth_resume_guard", False):
+        return
+
+    def _unsloth_train_with_resume_guard(self, *args, **kwargs):
+        resume_from_checkpoint = kwargs.get("resume_from_checkpoint", None)
+        if resume_from_checkpoint is None and len(args) != 0:
+            resume_from_checkpoint = args[0]
+
+        if resume_from_checkpoint:
+            _maybe_prepare_vllm_for_resume(self)
+        return original_train(self, *args, **kwargs)
+    pass
+
+    _unsloth_train_with_resume_guard._unsloth_resume_guard = True
+    trainer_class.train = _unsloth_train_with_resume_guard
+pass
+
 def PatchRL(FastLanguageModel):
     try:
         from trl.models.utils import unwrap_model_for_generation
@@ -1552,7 +1608,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         imports,
         overwrite = False,
     )
-
+    patched_trainer = getattr(created_module, f"Unsloth{RLTrainer_name}")
+    if trainer_file == "grpo_trainer":
+        _patch_resume_from_checkpoint_memory(patched_trainer)
     # Patch Trainer
     exec(
         f"trl.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}",
@@ -1784,6 +1842,16 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
             )
 
         init = init.replace(vllm_part, new_vllm_part)
+        else:
+            new_vllm_part = (
+                f"\n{' '*8}if {args}.use_vllm:\n"
+                f"{' '*12}self.llm = model.vllm_engine\n"
+                f"{' '*12}self.guided_decoding_regex = getattr(args, 'vllm_guided_decoding_regex', None)\n"
+                f"{' '*12}self._last_loaded_step = 0\n"
+                f"{' '*12}self.accelerator.wait_for_everyone()\n"
+                f"\n{' '*8}else:\n"
+            )
+            init = init.replace(vllm_part, new_vllm_part)
 
     # Search for vLLM calling in all child functions
     functions = dir(RLTrainer)

unsloth/models/rl_replacements.py

@@ -958,6 +958,10 @@ RL_PRE_ITEMS["grpo_trainer"].append(
 def grpo_trainer_compute_loss(function_name, function):
     if function_name != "compute_loss":
         return function
+    if "_compute_loss(" in function:
+        return function
+    if "_get_per_token_logps_and_entropies" in function:
+        return function
 
     def compute_loss(
         self, model, inputs, return_outputs = False, num_items_in_batch = None
@@ -1285,8 +1289,10 @@ RL_FUNCTIONS["kto_trainer"].append(kto_trainer_get_batch_logps)
 # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356
 # TRL warns if batch size is not a multiple of num_generations -> fix this.
 def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source):
-    if "divisible by the number of generations" not in RLTrainer_source:
-        # in later trl versions this doesn't exist anymore
+    if (
+        "divisible by the number of generations" not in RLTrainer_source
+        and "generation_batch_size" not in RLTrainer_source
+    ):
         return ""
     if "num_generations" not in RLConfig_source:
         return ""
@@ -1304,6 +1310,18 @@ def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source):
 RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size)
 
 
+def grpo_trainer_fix_generation_batch_size(RLTrainer_source, RLConfig_source):
+    if "generation_batch_size" not in RLConfig_source: return ""
+    if "steps_per_generation" not in RLConfig_source: return ""
+
+    check_generation_batch_size = \
+    "if generation_batch_size is not None and steps_per_generation is not None:\n"\
+    "    generation_batch_size = None\n"
+    return check_generation_batch_size
+pass
+RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_generation_batch_size)
+
+
 # Add other reward function names
 def grpo_trainer_metrics(RLTrainer_source, RLConfig_source):
     if "reward_funcs" not in RLTrainer_source:

unsloth/trainer.py

@@ -257,6 +257,10 @@ def _backwards_compatible_trainer(trainer_class, config_class):
 
         if ("args" in kwargs) and (Version(trl) >= Version("0.13.0.dev0")):
             training_args = kwargs.pop("args", None)
+            if isinstance(training_args, config_class):
+                kwargs["args"] = training_args
+                original_init(self, *args, **kwargs)
+                return
 
             # Get parameters that Trainer.__init__ actually expects
             trainer_params.remove("self")