From 74eaf52600d6f2d96c88c79e944fba4d2640a906 Mon Sep 17 00:00:00 2001 From: "abhishek.sharma" Date: Sat, 20 Dec 2025 11:47:03 +0530 Subject: [PATCH] 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