From 3f68dd5f0e9ecc54e75673cbce281d3ef356f35a Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Wed, 3 Jun 2026 18:44:54 +0530
Subject: [PATCH] Patch sibling config module so GRPOConfig resolves to the
 patched class (#5946)

Fixes #3931. After patching a TRL trainer, also patch the sibling config module (e.g. trl.trainer.grpo_config.GRPOConfig) to the Unsloth-patched config, so importing the config from its own module returns the patched class carrying unsloth_grpo_mini_batch. Defensive (try/except + hasattr) so it safely no-ops when no sibling config module exists.
---
 .../test_patch_trl_rl_trainers_defensive.py      | 16 ++++++++++++++++
 unsloth/models/rl.py                             | 12 ++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/tests/python/test_patch_trl_rl_trainers_defensive.py b/tests/python/test_patch_trl_rl_trainers_defensive.py
index 7c76ac2792..55a3425c44 100644
--- a/tests/python/test_patch_trl_rl_trainers_defensive.py
+++ b/tests/python/test_patch_trl_rl_trainers_defensive.py
@@ -67,3 +67,19 @@ def test_wrapper_swallows_impl_exception(monkeypatch):
 
     monkeypatch.setattr(_rl, "_patch_trl_rl_trainers_impl", _boom)
     assert _rl._patch_trl_rl_trainers("sft_trainer") is None
+
+
+def test_grpo_config_sibling_module_import_is_patched(tmp_path):
+    import unsloth  # noqa: F401
+    from trl import GRPOConfig as top_config
+    from trl.trainer import GRPOConfig as trainer_config
+    from trl.trainer.grpo_config import GRPOConfig as config_module_config
+    from trl.trainer.grpo_trainer import GRPOConfig as trainer_module_config
+
+    assert top_config is trainer_config
+    assert top_config is trainer_module_config
+    assert top_config is config_module_config
+
+    args = config_module_config(output_dir = str(tmp_path))
+    assert hasattr(args, "unsloth_grpo_mini_batch")
+    assert args.unsloth_grpo_mini_batch is None
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 359716fda4..afa2753994 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -19,6 +19,7 @@ __all__ = [
 
 import torch
 from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+import importlib
 import inspect
 import os
 import re
@@ -1845,6 +1846,17 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
         locals(),
         globals(),
     )
+    try:
+        config_module_name = trainer_file.replace("_trainer", "_config")
+        config_module = importlib.import_module(f"trl.trainer.{config_module_name}")
+        if hasattr(config_module, RLConfig_name):
+            setattr(
+                config_module,
+                RLConfig_name,
+                getattr(created_module, f"Unsloth{RLConfig_name}"),
+            )
+    except Exception:
+        pass
 
     if trainer_file == "grpo_trainer":
         try: