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.
This commit is contained in:
parent
aa0db1ff5b
commit
3f68dd5f0e
2 changed files with 28 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue