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: