diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index d91a6680d4..24a5c8d1f9 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -25,6 +25,7 @@ from unsloth_zoo.compiler import create_new_function from unsloth_zoo.logging_utils import PatchRLStatistics from .rl_replacements import ( RL_EXTRA_ARGS, + RL_FUNCTIONS, ) def PatchRL(FastLanguageModel): @@ -365,8 +366,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): RLConfig_extra_args = extra_args RLConfig_call_args = call_args - # Patch vLLM - RLTrainer_extras = patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports) + # Patch vLLM and other functions + RLTrainer_extras = patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports) if RLTrainer_extras is None: RLTrainer_extras = f"_Unsloth{RLTrainer_name} = {RLTrainer_name}" @@ -414,7 +415,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): pass -def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): +def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): init = inspect.getsource(RLTrainer.__init__) old_init = init @@ -475,6 +476,7 @@ def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): functions = [x for x in functions if f"def {x}" in RLTrainer_source] changed = {"__init__" : (old_init, init,)} + edit_functions = RL_FUNCTIONS.get(trainer_file, []) for function in functions: if not hasattr(RLTrainer, function): continue @@ -483,6 +485,11 @@ def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): except: continue original_source = source + # Check for function + for edit_function in edit_functions: + source = edit_function(function, source) + pass + # 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", diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index a09fcb1fb2..56c5c7ad9e 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -14,14 +14,19 @@ __all__ = [ "RL_EXTRA_ARGS", + "RL_FUNCTIONS", ] -RL_EXTRA_ARGS = dict() +import re +from collections import defaultdict +RL_EXTRA_ARGS = defaultdict(list) +RL_FUNCTIONS = defaultdict(list) + def sft_trainer_fix_untraiend_tokens(call_args, extra_args): if "model" in call_args and "train_dataset" in call_args: fix_tokenizer = \ - "IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', set())\n"\ + "IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', '').split('\n')\n"\ "from unsloth_zoo.tokenizer_utils import fix_untrained_tokens\n"\ "from unsloth_zoo.training_utils import fix_zero_training_loss\n"\ "if 'tokenizer' not in locals(): tokenizer = processing_class\n"\ @@ -30,7 +35,7 @@ def sft_trainer_fix_untraiend_tokens(call_args, extra_args): return fix_tokenizer return "" pass -RL_EXTRA_ARGS["sft_trainer"] = [sft_trainer_fix_untraiend_tokens,] +RL_EXTRA_ARGS["sft_trainer"].append(sft_trainer_fix_untraiend_tokens) def dpo_trainer_fix_columns(call_args, extra_args): @@ -47,4 +52,45 @@ def dpo_trainer_fix_columns(call_args, extra_args): return fix_dpo return "" pass -RL_EXTRA_ARGS["dpo_trainer"] = [dpo_trainer_fix_columns,] +RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_columns) + + +def sft_trainer_prepare_dataset(function_name, function): + if function_name != "_prepare_non_packed_dataloader" and \ + function_name != "_prepare_dataset": return + + check_text = \ + "\n"\ + "if 'tokenizer' not in locals(): tokenizer = processing_class\n"\ + "if 'formatting_func' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `formatting_func` does not exist!')\n"\ + "if 'dataset_text_field' not in locals() and 'args' in locals(): dataset_text_field = args.dataset_text_field\n"\ + "if 'dataset_text_field' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `dataset_text_field` does not exist!')\n"\ + "test_text = dataset[0][dataset_text_field] if (formatting_func is None and dataset_text_field is not None) else formatting_func(dataset[0])[0]\n"\ + "chat_template = getattr(tokenizer, 'chat_template', None)\n"\ + "chat_template = '' if chat_template is None else chat_template\n"\ + "has_bos_token_already = (test_text.startswith(tokenizer.bos_token) or tokenizer.bos_token in chat_template) "\ + "if getattr(tokenizer, 'bos_token', None) is not None else False\n"\ + "if 'add_special_tokens' not in locals() and has_bos_token_already:\n"\ + " from functools import partial\n"\ + " tokenizer = partial(tokenizer, add_special_tokens = False)\n"\ + " processing_class = tokenizer\n"\ + "else:\n"\ + " add_special_tokens = False if has_bos_token_already else add_special_tokens\n" + + check_text = check_text.split("\n") + check_text = "\n".join(" "*where + x for x in check_text) + check_text = check_text.rstrip() + "\n" + + # .*? matches first match. .+? matches final match. + replacer = re.findall( + f"def {function_name}\(.*?\).*?\:\n", + function, + flags = re.MULTILINE | re.DOTALL, + ) + if len(replacer) != 0: + replacer = replacer[0] + function = function.replace(replacer, replacer + check_text) + pass + return function +pass +RL_FUNCTIONS["sft_trainer"].append(sft_trainer_prepare_dataset)