diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index ca1a1b5dbd..b653fb960d 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -39,6 +39,7 @@ import os import re import functools + def PatchRL(FastLanguageModel): from trl.models.utils import unwrap_model_for_generation @@ -240,6 +241,7 @@ pass def PatchRLStatistics(algorithm = "GRPO"): + # Get notebook statistics columns to show up algorithm = algorithm.upper() all_metrics = get_trl_metrics() if algorithm not in all_metrics: @@ -252,7 +254,137 @@ def PatchRLStatistics(algorithm = "GRPO"): pass +def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): + # Patch for vLLM and Unsloth PEFT + import trl.trainer + + trainer = eval(f"trl.trainer.{trainer_file}") + name = [x for x in dir(trainer) if x.endswith("Trainer") and x != "Trainer" and trainer_file.split("_")[0] in x.lower()] + assert(len(name) == 1) + RLTrainer_name = name[0] + RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}") + + try: + __init__ = inspect.getsource(RLTrainer.__init__) + except: + # Already patched most likely! + return + all_imports = dir(trainer) + imports = [x for x in all_imports if not x.startswith("_") and x in __init__] + + spaces = __init__.find("def") + __init__ = __init__.split("\n") + __init__ = "\n".join(x[spaces:] for x in __init__) + + vllm_part = re.findall( + r"(\n[\s]{4}"\ + r"if (self|args)\.use_vllm\:.+?"\ + r"\n[\s]{4,}"\ + "else:\n)", + __init__, + flags = re.MULTILINE | re.DOTALL, + ) + if (len(vllm_part) != 1): return + + vllm_part, args = vllm_part[0][0], vllm_part[0][1] + # Strip all comments + new_vllm_part = re.sub(r"\#[^\n]{1,}\n", "", vllm_part) + + # Get SamplingParams + sampling_params = re.findall( + r"\n[\s]{4,}(self\.[^\s]{1,}[\s]{0,}\=[\s]{0,}"\ + r"SamplingParams\(.+?\))", + new_vllm_part, + flags = re.MULTILINE | re.DOTALL, + ) + if len(sampling_params) != 1: return + + sampling_params = sampling_params[0] + sampling_params = \ + " "*8 + "self.llm = model.vllm_engine; " + \ + sampling_params # Add spaces + new_vllm_part = f"\n if {args}.use_vllm:\n{sampling_params}\n else:\n" + __init__ = __init__.replace(vllm_part, new_vllm_part) + + # Remove peft_config + __init__ = __init__.replace("elif peft_config is None:", "elif False:") + __init__ = __init__.replace("elif peft_config is not None:", "elif False:") + __init__ = __init__.replace("if peft_config is None:", "if False:") + __init__ = __init__.replace("if peft_config is not None:", "if False:") + __init__ = __init__.replace("get_peft_model(model, peft_config)", "model") + + # Search for vLLM calling in all child functions + functions = dir(RLTrainer) + RLTrainer_source = inspect.getsource(RLTrainer) + functions = [x for x in functions if f"def {x}" in RLTrainer_source] + + changed = {"__init__" : __init__} + for function in functions: + if not hasattr(RLTrainer, function): continue + fx = getattr(RLTrainer, function) + try: + source = inspect.getsource(fx) + except: + continue + original_source = source + + # llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + source = re.sub( + r"(\n[\s]{4,}).+?model_executor\.driver_worker.+?\n", + r"\n\1pass\n", + source, + ) + # llm_model.load_weights(model.state_dict().items()) + source = re.sub( + r"(\n[\s]{4,}).+?load_weights\(.+?\n", + r"\n\1pass\n", + source, + ) + # Replace self.llm.generate and self.llm.chat + lora_name = trainer_file + "_lora_model" + source = re.sub( + r"(self\.llm\.(?:generate|chat)\([^\)]{1,})\)", + r"\1, lora_request = model.load_lora('" + lora_name + r"', load_tensors = True))", + source + ) + if source == original_source: continue + + # Find all imports + imports += [x for x in all_imports if not x.startswith("_") and x in source] + + # Create actual function + spaces = source.find("def") + source = source.split("\n") + source = "\n".join(x[spaces:] for x in source) + changed[function] = source + pass + + # Import all functions + imports = list(set(imports)) + imports = f"from trl.trainer.{trainer_file} import (\n" + ',\n'.join(imports) + ")" + exec(imports) + + # Patch all functions + for function in changed: + exec(changed[function]) + exec(f"trl.trainer.{trainer_file}.{RLTrainer_name}.{function} = {function}") + pass +pass + + +def patch_trl_rl_trainers(): + # Patch all TRL modules if they have vLLM or PEFT + import trl.trainer + all_trainers = dir(trl.trainer) + all_trainers = [x for x in all_trainers if x.islower() and x.endswith("_trainer")] + for trainer in all_trainers: + _patch_trl_rl_trainers(trainer) + return +pass + + def PatchFastRL(algorithm = "GRPO", FastLanguageModel = None): if FastLanguageModel is not None: PatchRL(FastLanguageModel) + patch_trl_rl_trainers() PatchRLStatistics(algorithm) pass