From 0553244b9b259e5399db2b018576103ebd0740c5 Mon Sep 17 00:00:00 2001 From: yash solanki Date: Tue, 10 Mar 2026 20:42:53 +0530 Subject: [PATCH] Fix GRPO resume on newer TRL/vLLM stack --- unsloth/kernels/fast_lora.py | 10 +++++ unsloth/models/llama.py | 2 + unsloth/models/rl.py | 70 ++++++++++++++++++++++++++++++- unsloth/models/rl_replacements.py | 22 +++++++++- unsloth/trainer.py | 4 ++ 5 files changed, 105 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index f1c0e298d9..c8ba88c72d 100644 --- a/unsloth/kernels/fast_lora.py +++ b/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()) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 93d93e26d6..b1fdba577d 100644 --- a/unsloth/models/llama.py +++ b/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, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index b98545fcd2..c288bf7060 100755 --- a/unsloth/models/rl.py +++ b/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) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 14153d30b4..476286e0ba 100755 --- a/unsloth/models/rl_replacements.py +++ b/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: diff --git a/unsloth/trainer.py b/unsloth/trainer.py index b41e94fea2..3734702731 100644 --- a/unsloth/trainer.py +++ b/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")