From 6329e4bf37c6af91f3c8582e9a8e2ca145a3db10 Mon Sep 17 00:00:00 2001
From: yash solanki <alphacr792@gmail.com>
Date: Tue, 10 Mar 2026 23:43:28 +0530
Subject: [PATCH] Fix review feedback for trainer config remapping

---
 unsloth/models/_utils.py          |  2 +-
 unsloth/models/rl.py              |  6 ++++++
 unsloth/models/rl_replacements.py | 19 +++++++++++++++++++
 unsloth/trainer.py                |  8 +++++++-
 4 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 19b5fe0574..2746a70e45 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -730,7 +730,7 @@ try:
     # Some Config files use layer_type_validation
     # for eg Gemma-2, so we must import it to stop errors.
     from transformers.configuration_utils import layer_type_validation
-except:
+except ImportError:
     pass
 
 try:
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index d0bddd625e..350c8f9d41 100755
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -102,6 +102,7 @@ def _maybe_prepare_vllm_for_resume(trainer):
         llm = getattr(getattr(trainer, "model", None), "vllm_engine", None)
 
     sleep_fn = getattr(llm, "sleep", None)
+    slept = False
     if callable(sleep_fn):
         try:
             sleep_mode = int(os.environ.get("VLLM_SLEEP_MODE", "1"))
@@ -116,8 +117,10 @@ def _maybe_prepare_vllm_for_resume(trainer):
         try:
             if signature is not None and len(signature.parameters) == 0:
                 sleep_fn()
+                slept = True
             else:
                 sleep_fn(sleep_mode)
+                slept = True
         except Exception as error:
             logger.warning_once(
                 f"Unsloth: vLLM sleep() failed during resume cleanup: {error}"
@@ -129,6 +132,9 @@ def _maybe_prepare_vllm_for_resume(trainer):
         gc.collect()
         torch.cuda.empty_cache()
 
+    if slept:
+        trainer._unsloth_resume_wake_vllm = True
+
 
 def _patch_resume_from_checkpoint_memory(trainer_class):
     original_train = getattr(trainer_class, "train", None)
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 314feb5d2a..14153d30b4 100755
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -236,6 +236,25 @@ def grpo_trainer__prepare_inputs(function_name, function):
     if function_name != "_prepare_inputs":
         return function
 
+    import re
+
+    match = re.search(
+        r"^([ \t]*)with torch\.inference_mode\(\):", function, flags = re.MULTILINE
+    )
+    if match is not None:
+        indent = match.group(1)
+        nested_indent = indent + " " * 4
+        wake_block = (
+            f"{indent}if getattr(self, '_unsloth_resume_wake_vllm', False):\n"
+            f"{nested_indent}wake_up = getattr(getattr(self, 'llm', None), 'wake_up', None)\n"
+            f"{nested_indent}if callable(wake_up):\n"
+            f"{nested_indent}    try: wake_up()\n"
+            f"{nested_indent}    except Exception as error:\n"
+            f"{nested_indent}        logging.getLogger(__name__).warning('Unsloth: vLLM wake_up() failed during resume: %s', error)\n"
+            f"{nested_indent}self._unsloth_resume_wake_vllm = False\n"
+            f"{indent}with torch.inference_mode():"
+        )
+        function = function[: match.start()] + wake_block + function[match.end() :]
     # Add mixed precision training
     function = function.replace(
         "with torch.inference_mode():",
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 65abe6801f..b41e94fea2 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -303,7 +303,13 @@ def _backwards_compatible_trainer(trainer_class, config_class):
             # causes the 2nd init to fail as there are mutual exclusive checks on pairs of parameters.
             # Refer: https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_config.py#L499-L502 for example
             # So we only create config class if the previous init was not TrainingArguments
-            if not isinstance(training_args, TrainingArguments):
+            if isinstance(training_args, config_class):
+                import copy
+
+                config = copy.deepcopy(training_args)
+                for key, value in additional_config_kwargs.items():
+                    setattr(config, key, value)
+            elif not isinstance(training_args, TrainingArguments):
                 config = config_class(**config_dict)
             else:
                 config = training_args