From 74eaf52600d6f2d96c88c79e944fba4d2640a906 Mon Sep 17 00:00:00 2001 From: "abhishek.sharma" Date: Sat, 20 Dec 2025 11:47:03 +0530 Subject: [PATCH 1/9] Fix model training state restoration in GRPO trainer Store the model's training state before generation and restore inference mode after completion if the model wasn't originally in training mode. This ensures the model returns to the correct state after generate and score operations. --- unsloth/models/rl_replacements.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 7d4d520c1f..dd139ffd25 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -259,6 +259,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # The new multi-line string that will replace the line above replacement_lines = """ batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + _was_training = self.model.training try: # TRL 0.23.1 and below path if not has_images: @@ -387,6 +388,13 @@ def grpo_trainer__generate_and_score_completions(function_name, function): patched = patched[: match.start()] + wrapped + patched[match.end() :] function = patched + + function = function.replace( + " return output", # 8 spaces before 'return' + """ if not _was_training: + self.model.for_inference() + return output""" + ) return function From 8b5130ae2d85a827a6d37e0dbfb9dabea6c1e05e Mon Sep 17 00:00:00 2001 From: numb3r33 Date: Sat, 20 Dec 2025 12:30:33 +0530 Subject: [PATCH 2/9] Remove the comment. --- unsloth/models/rl_replacements.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index dd139ffd25..248c5aab85 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -390,7 +390,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function): function = patched function = function.replace( - " return output", # 8 spaces before 'return' + " return output", """ if not _was_training: self.model.for_inference() return output""" From 806f8d2d7e34712aff966e302cd8baff0448e9bc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Dec 2025 07:02:30 +0000 Subject: [PATCH 3/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/rl_replacements.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 248c5aab85..e13e5e6d78 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -388,12 +388,12 @@ def grpo_trainer__generate_and_score_completions(function_name, function): patched = patched[: match.start()] + wrapped + patched[match.end() :] function = patched - + function = function.replace( - " return output", - """ if not _was_training: + " return output", + """ if not _was_training: self.model.for_inference() - return output""" + return output""", ) return function From ab81842f78a07307ecc77358a53d5621dc1220f2 Mon Sep 17 00:00:00 2001 From: numb3r33 Date: Wed, 24 Dec 2025 00:14:50 +0530 Subject: [PATCH 4/9] Fix indentation handling in grpo_trainer return statement replacement Use regex to dynamically detect and preserve the original indentation when replacing the 'return output' statement, instead of hardcoding spaces. This ensures the patched code maintains consistent indentation regardless of the original formatting. --- unsloth/models/rl_replacements.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index e13e5e6d78..bcac699b3f 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -388,13 +388,17 @@ def grpo_trainer__generate_and_score_completions(function_name, function): patched = patched[: match.start()] + wrapped + patched[match.end() :] function = patched + + match = re.search(r'^(\s*)return output', function, re.MULTILINE) - function = function.replace( - " return output", - """ if not _was_training: - self.model.for_inference() - return output""", - ) + if match: + indent = match.group(1) + function = function.replace( + f"{indent}return output", + f"""{indent}if not _was_training: + {indent} self.model.for_inference() + {indent}return output""" + ) return function From c8784ec87e70df30936731bf141a0ab33f2eda50 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 23 Dec 2025 19:20:08 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/rl_replacements.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index bcac699b3f..5158019132 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -388,8 +388,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): patched = patched[: match.start()] + wrapped + patched[match.end() :] function = patched - - match = re.search(r'^(\s*)return output', function, re.MULTILINE) + + match = re.search(r"^(\s*)return output", function, re.MULTILINE) if match: indent = match.group(1) @@ -397,7 +397,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function): f"{indent}return output", f"""{indent}if not _was_training: {indent} self.model.for_inference() - {indent}return output""" + {indent}return output""", ) return function From e4f582bdc25ed1076c3e107698bdc2a4d7e1feae Mon Sep 17 00:00:00 2001 From: numb3r33 Date: Wed, 24 Dec 2025 01:02:03 +0530 Subject: [PATCH 6/9] Refactor return statement replacement to use explicit newlines Replace f-string triple-quoted approach with explicit newline characters for clearer string construction in the grpo_trainer patch. --- unsloth/models/rl_replacements.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 5158019132..8436ca0dc9 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -393,12 +393,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): if match: indent = match.group(1) - function = function.replace( - f"{indent}return output", - f"""{indent}if not _was_training: - {indent} self.model.for_inference() - {indent}return output""", - ) + new_code = indent + "if not _was_training:\n" + indent + " self.model.for_inference()\n" + indent + "return output" + function = function.replace(f"{indent}return output", new_code) return function From ef3e2b39a8a9aacae67f4b24e7183fcd04efcd4f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 23 Dec 2025 20:06:36 +0000 Subject: [PATCH 7/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/rl_replacements.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 8436ca0dc9..f0f0386bd1 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -393,7 +393,14 @@ def grpo_trainer__generate_and_score_completions(function_name, function): if match: indent = match.group(1) - new_code = indent + "if not _was_training:\n" + indent + " self.model.for_inference()\n" + indent + "return output" + new_code = ( + indent + + "if not _was_training:\n" + + indent + + " self.model.for_inference()\n" + + indent + + "return output" + ) function = function.replace(f"{indent}return output", new_code) return function From 6918e2d31a6caee5ce75ce9a20e1e67908561b28 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 5 Jan 2026 13:50:48 +0000 Subject: [PATCH 8/9] Fix GRPO training state restoration --- unsloth/models/rl.py | 46 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 4ea36519d9..1327208c46 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -238,12 +238,18 @@ def prepare_for_training_mode(f): @functools.wraps(f) def wrapper(self, *args, **kwargs): # Enable training mode + _was_training = None + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): self.model.for_training() output = f(self, *args, **kwargs) - # Return inference mode + # Restore previous mode when possible if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - self.model.for_inference() + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training() # Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run try: import wandb @@ -323,6 +329,32 @@ pass ''' +def _wrap_grpo_generate_and_score(trainer_cls): + if not hasattr(trainer_cls, "_generate_and_score_completions"): + return + original = trainer_cls._generate_and_score_completions + if getattr(original, "_unsloth_restore_training_wrapped", False): + return + + def wrapped(self, *args, **kwargs): + was_training = getattr(getattr(self, "model", None), "training", None) + try: + return original(self, *args, **kwargs) + finally: + if ( + was_training is False + and hasattr(self, "model") + and hasattr(self.model, "for_inference") + ): + try: + self.model.for_inference() + except Exception: + pass + + wrapped._unsloth_restore_training_wrapped = True + trainer_cls._generate_and_score_completions = wrapped + + def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Patch for vLLM and Unsloth PEFT import trl @@ -1046,6 +1078,16 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): globals(), ) + if trainer_file == "grpo_trainer": + try: + _wrap_grpo_generate_and_score( + getattr(created_module, f"Unsloth{RLTrainer_name}") + ) + except Exception as e: + logger.info( + f"Unsloth: Could not wrap _generate_and_score_completions for {RLTrainer_name}: {e}" + ) + def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): init = inspect.getsource(RLTrainer.__init__) From ef533cddf70d258345ea45212f59b0703672ae20 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 5 Jan 2026 13:55:08 +0000 Subject: [PATCH 9/9] Revert rl_replacements GRPO edits --- unsloth/models/rl_replacements.py | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index f0f0386bd1..5e079335ae 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -259,7 +259,6 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # The new multi-line string that will replace the line above replacement_lines = """ batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size - _was_training = self.model.training try: # TRL 0.23.1 and below path if not has_images: @@ -389,20 +388,6 @@ def grpo_trainer__generate_and_score_completions(function_name, function): function = patched - match = re.search(r"^(\s*)return output", function, re.MULTILINE) - - if match: - indent = match.group(1) - new_code = ( - indent - + "if not _was_training:\n" - + indent - + " self.model.for_inference()\n" - + indent - + "return output" - ) - function = function.replace(f"{indent}return output", new_code) - return function @@ -876,13 +861,19 @@ def grpo_trainer_compute_loss(function_name, function): else torch.tensor(0.0, device = self.model.device) ) self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( - nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + self.accelerator.gather(min_importance_sampling_ratio) + .nan_to_num(nan = float("inf")) + .min() + .item() ) self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() ) self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( - nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + self.accelerator.gather(max_importance_sampling_ratio) + .nan_to_num(nan = float("-inf")) + .max() + .item() ) return loss @@ -964,11 +955,15 @@ def openenv_vllm_reload_weights(): return if Version(importlib_version("trl")) < Version("0.26.0"): return + try: import trl.experimental.openenv.utils as openenv_utils import trl.experimental.openenv as openenv except ImportError as e: logger.info(f"Unsloth: Failed to import trl openenv: {e}") + logger.info( + "Unsloth: trl.experimental.openenv not available — skipping RL openenv patches." + ) return src = inspect.getsource(openenv_utils.generate_rollout_completions)