diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 515c6587f7..5d6117b704 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -16,29 +16,13 @@ __all__ = [ "PatchFastRL", ] -METRICS_MOVE_TO_END = [ - "nll", - "aux", - "beta", - "alpha", -] import torch -try: - from transformers.utils.notebook import ( - IntervalStrategy, - NotebookTrainingTracker, - NotebookProgressCallback, - ) - HAS_NOTEBOOK = True -except: - HAS_NOTEBOOK = False -pass from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import inspect import os import re -import functools from unsloth_zoo.compiler import create_new_function +from unsloth_zoo.logging_utils import PatchRLStatistics def PatchRL(FastLanguageModel): @@ -78,219 +62,290 @@ def PatchRL(FastLanguageModel): trainers = [x for x in trainers if x.endswith("_trainer")] unwrap = "unwrap_model_for_generation" for trainer in trainers: - if hasattr(eval(f"trl.trainer.{trainer}"), unwrap): - exec(f"trl.trainer.{trainer}.{unwrap} = unsloth_{unwrap}") + try: current_trainer = eval(f"trl.trainer.{trainer}") + except: continue + if hasattr(current_trainer, unwrap): + try: exec(f"trl.trainer.{trainer}.{unwrap} = unsloth_{unwrap}") + except: continue pass pass -def NotebookProgressCallback_on_train_begin(Trainer_metrics): - def _NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs): - self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step" - self.training_loss = 0 - self.last_log = 0 - column_names = [self.first_column] + ["Training Loss"] - if args.eval_strategy != IntervalStrategy.NO: - column_names.append("Validation Loss") - column_names += [x.replace("/", " / ") for x in Trainer_metrics] - self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names) - pass - return _NotebookProgressCallback_on_train_begin +RLTrainer_replacement = ''' +from typing import * +from dataclasses import dataclass, field + +@dataclass +class Unsloth{RLConfig_name}({RLConfig_name}): + """ + {__RLConfig_doc__} + """ + sampling_params: Optional[Any] = field( + default = None, + metadata = {{'help': 'vLLM SamplingParams'}}, + ) + def __init__({RLConfig_arguments}, + sampling_params = None + ): +{RLConfig_extra_args} + super().__init__({RLConfig_call_args}) pass +{RLTrainer_extras} -def NotebookProgressCallback_on_log(Trainer_metrics): - def _NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwargs): - # Only for when there is no evaluation - if args.eval_strategy == IntervalStrategy.NO and "loss" in logs: - values = {"Training Loss": logs["loss"]} - for metric in Trainer_metrics: - # Sometimes metric is not inside logs - try: values[metric.replace("/", " / ")] = logs[metric] - except: pass - pass - # First column is necessarily Step since we're not in epoch eval strategy - values["Step"] = state.global_step - self.training_tracker.write_line(values) - pass - pass - return _NotebookProgressCallback_on_log +class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): + """ + {__RLTrainer_doc__} + """ + def __init__({RLTrainer_arguments} + ): + if args is None: args = Unsloth{RLConfig_name}() +{RLTrainer_extra_args} + super().__init__({RLTrainer_call_args}) pass - - -def NotebookTrainingTracker_write_line(Trainer_metrics): - set_Trainer_metrics = set(Trainer_metrics) - def _NotebookTrainingTracker_write_line(self, values): - """ - Write the values in the inner table. - - Args: - values (`Dict[str, float]`): The values to display. - """ - if self.inner_table is None: - self.inner_table = [list(values.keys()), list(values.values())] - else: - columns = self.inner_table[0] - new_values = {} - for key, value in values.items(): - lowered = key.lower() - if lowered in set_Trainer_metrics: - new_values[lowered.replace("/", " / ")] = value - else: - new_values[key] = value - pass - values = new_values - - self.inner_table[0] = columns - if len(self.inner_table) > 1: - last_values = self.inner_table[-1] - first_column = self.inner_table[0][0] - if last_values[0] != values[first_column]: - # write new line - self.inner_table.append([values[c] if c in values else "No Log" for c in columns]) - else: - # update last line - new_values = values - for c in columns: - if c not in new_values.keys(): - new_values[c] = last_values[columns.index(c)] - self.inner_table[-1] = [new_values[c] for c in columns] - else: - # Edit for evaluation purposes - self.inner_table.append([values[c] if c in values else 0 for c in columns]) - pass - pass - pass - return _NotebookTrainingTracker_write_line -pass - - -def _PatchRLStatistics(metrics, algorithm): - if HAS_NOTEBOOK: - if len(metrics) == 0: - raise RuntimeError(f"Unsloth: RL statistics for {algorithm} failed with no metrics seen?") - from transformers.trainer import is_in_notebook - if is_in_notebook(): - # Patch DPO notebook printing - NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line(metrics) - from transformers.trainer import DEFAULT_PROGRESS_CALLBACK - DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin(metrics) - DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log(metrics) - pass - pass -pass - - -@functools.cache -def get_trl_metrics(): - # Gets metrics so we can output them in notebooks - - import trl.trainer - trainers = dir(trl.trainer) - trainers = [x for x in trainers if x.endswith("_trainer")] - filepath = inspect.getfile(trl.trainer) - filepath = os.path.split(filepath)[0] - - all_metrics = dict() - for trainer in trainers: - filename = os.path.join(filepath, f"{trainer}.py") - if not os.path.exists(filename): continue - with open(filename, "r") as file: file = file.read() - - # Get metrics['kl'] or stats['kl'] - metrics = re.findall(r"metrics\[[\"\']([^\"\']{1,})[\"\']\]", file) - stats = re.findall(r"stats\[[\"\']([^\"\']{1,})[\"\']\]", file) - metrics = metrics + stats - - # Get optional f-strings - metrics_f = re.findall(r"metrics\[f[\"\']\{[^\}]{1,}\}([^\"\']{1,})[\"\']\]", file) - stats_f = re.findall(r"stats\[f[\"\']\{[^\}]{1,}\}([^\"\']{1,})[\"\']\]", file) - metrics_f = metrics_f + stats_f - # Filter out prefixes if seen - # metrics[f"{prefix}rewards/chosen"] - left_prefix = 'prefix = "eval_" if train_eval == "eval" else ""' in file - if left_prefix: metrics += metrics_f - - # Move all eval_ things to the end and reward to the front - beginning = [] - middle = [] - end = [] - for x in metrics: - lowered = x.lower() - if "reward" in lowered: - beginning.append(x) - elif x.lower().startswith("eval"): - end.append(x) - else: - # Check if we want to move to the end - moved = False - for move_end in METRICS_MOVE_TO_END: - if move_end in lowered: - end.append(x) - moved = True - break - if not moved: - middle.append(x) - pass - pass - metrics = beginning + middle + end - - all_metrics[trainer[:trainer.find("_")].upper()] = metrics - pass - return all_metrics -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: - print( - f"Unsloth for {algorithm.upper()} is not yet implemented! Just ignore this function.\n"\ - f"We support: `{list(all_metrics.keys())}`" - ) - pass - _PatchRLStatistics(all_metrics[algorithm], algorithm) -pass - +''' def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Patch for vLLM and Unsloth PEFT import trl 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! + trainer = eval(f"trl.trainer.{trainer_file}") + except Exception as error: return - old__init__ = __init__ + + # Get SFTTrainer and SFTConfig names + name = [x for x in dir(trainer) if x.endswith("Trainer") and x != "Trainer" and trainer_file.split("_")[0] in x.lower()] + config = [x for x in dir(trainer) if x.endswith("Config") and x != "Config" and trainer_file.split("_")[0] in x.lower()] + if len(name) != 1: return + if len(config) != 1: return + + # Get SFTTrainer, SFTConfig + RLTrainer_name = name[0] + RLConfig_name = config[0] + try: RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}") + except: return + try: RLConfig = eval(f"trl.trainer.{trainer_file}.{RLConfig_name}" ) + except: return + + # Check name + if RLTrainer.__name__.startswith("Unsloth"): return + if RLConfig .__name__.startswith("Unsloth"): return + all_imports = dir(trainer) - assert("Union" in all_imports) imports = [x for x in all_imports if not x.startswith("_")] - imports += ["Trainer"] - spaces = __init__.find("def") - __init__ = __init__.split("\n") - __init__ = "\n".join(x[spaces:] for x in __init__) + # Get default arguments + EMPTY = inspect.Parameter.empty + processed = [] + for RLobject in [RLTrainer, RLConfig]: + parameters = inspect.signature(RLobject.__init__).parameters + types = (bool, type(None), int, float, str,) + arguments = ["self"] + call_args = [] + for k, v in parameters.items(): + if k == "self": continue + v = v.default + if v == "\n": v = re.escape("\n") + if v is EMPTY: arguments.append(k) + elif type(v) is str: arguments.append(f"{k} = '{v}'") + elif type(v) in types: arguments.append(f"{k} = {v}") + else: continue + call_args.append(f"{k} = {k}") + pass + arguments = f"\n{' '*8}" + f",\n{' '*8}".join(arguments) + call_args = f"\n{' '*12}" + f",\n{' '*12}".join(call_args) + processed.append((arguments, call_args,)) + pass + + # Process RLTrainer first + arguments, call_args = processed[0] + + # Add tokenizer if not seen + if "tokenizer" not in parameters and "processing_class" in parameters: + arguments += f",\n{' '*8}tokenizer = None" + call_args = call_args.replace( + "processing_class = processing_class", + "processing_class = tokenizer if tokenizer is not None else processing_class", + ) + pass + + # Edit bf16, fp16 by checking model's torch_dtype directly + extra_args = "" + if "args" in call_args: + mixed_precision = \ + "use_bf16 = getattr(args, 'bf16', False)\n"\ + "use_fp16 = getattr(args, 'fp16', False)\n"\ + "dtype = getattr(model.config, 'torch_dtype', None)\n"\ + "if dtype is None: dtype = model.get_input_embeddings().dtype\n"\ + "from unsloth_zoo.utils import _get_dtype\n"\ + "dtype = _get_dtype(dtype)\n"\ + "float16 = dtype == torch.float16\n"\ + "if float16 and use_bf16: raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n"\ + "if not float16 and use_fp16: raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"\ + "if not use_bf16 and not use_fp16:\n"\ + " args.fp16 = float16\n"\ + " args.bf16 = not float16\n" + extra_args += mixed_precision + pass + + # Check if per_device_eval_batch_size (default 8) bigger than bsz + # Also use FP16 / BF16 evaluation + if "args" in call_args: + # Check eval_dataset first + if "eval_dataset" in call_args: + check_eval_dataset = \ + "if getattr(args, 'eval_strategy', 'no') == 'no':\n"\ + " args.eval_strategy = 'steps'\n"\ + " if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1\n" + extra_args += check_eval_dataset + pass + + eval_changes = \ + "ga_steps = getattr(args, 'gradient_accumulation_steps', None)\n"\ + "if getattr(args, 'eval_strategy', 'no') != 'no':\n"\ + " eval_bsz = getattr(args, 'per_device_eval_batch_size', 8)\n"\ + " if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size\n"\ + " if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps\n"\ + "fp16_full_eval = getattr(args, 'fp16_full_eval', False)\n"\ + "bf16_full_eval = getattr(args, 'bf16_full_eval', False)\n"\ + "if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True\n"\ + "if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False\n"\ + "if not bf16_full_eval and not fp16_full_eval: args.bf16_full_eval = args.bf16; args.fp16_full_eval = args.fp16\n" + + extra_args += eval_changes + pass + + # Add statistics as well! + extra_args += \ + "from unsloth_zoo.logging_utils import PatchRLStatistics\n"\ + f"PatchRLStatistics('{trainer_file}')\n" + + # Create RLTrainer args + extra_args = extra_args.split("\n") + extra_args = "\n".join(" "*8 + x for x in extra_args) + RLTrainer_arguments = arguments + RLTrainer_extra_args = extra_args + RLTrainer_call_args = call_args + + # Fix RLConfig next + arguments, call_args = processed[1] + extra_args = "" + + # Edit GA / bsz and weight_decay + replacements = { + "output_dir" : 'unsloth_training_checkpoints', + "logging_nan_inf_filter" : False, + "per_device_train_batch_size" : 4, + "gradient_accumulation_steps" : 2, + "weight_decay" : 0.01, + "warmup_ratio" : 0.1, + "seed" : 3407, + "optim" : "adamw_8bit", + "learning_rate" : 5e-05, + "per_device_eval_batch_size" : 4, + "eval_accumulation_steps" : 2, + "torch_empty_cache_steps" : 250, + } + for k, v in replacements.items(): + x = f"{k}( = [^,\n]{{1,}})?,\n" + y = f"'{v}'" if type(v) is str else f"{v}" + y = f"{k} = {y},\n" + arguments = re.sub(x, y, arguments) + pass + + # Warn on too large or too small learning rate + if " learning_rate" in call_args: + learning_rate_check = \ + "if learning_rate < 1e-7: raise FloatingPointError(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!')\n"\ + "if learning_rate > 1: raise OverflowError(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!')" + extra_args += learning_rate_check + pass + + # Create RLConfig args + extra_args = extra_args.split("\n") + extra_args = "\n".join(" "*8 + x for x in extra_args) + RLConfig_arguments = arguments + RLConfig_extra_args = extra_args + RLConfig_call_args = call_args + + # Patch vLLM + RLTrainer_extras = patch_vllm(trainer_file, RLTrainer_name, all_imports, imports) + if RLTrainer_extras is None: + RLTrainer_extras = f"_Unsloth{RLTrainer_name} = {RLTrainer_name}" + + # Create full module + exec(f"from trl.trainer import ({RLTrainer_name}, {RLConfig_name},)") + __RLTrainer_doc__ = eval(f"trl.trainer.{RLTrainer_name}").__doc__ + __RLConfig_doc__ = eval(f"trl.trainer.{RLConfig_name}") .__doc__ + + RLTrainer_source = RLTrainer_replacement.format( + RLTrainer_name = RLTrainer_name, + __RLTrainer_doc__ = __RLTrainer_doc__, + RLTrainer_arguments = RLTrainer_arguments, + RLTrainer_extra_args = RLTrainer_extra_args, + RLTrainer_call_args = RLTrainer_call_args, + + RLConfig_name = RLConfig_name, + __RLConfig_doc__ = __RLConfig_doc__, + RLConfig_arguments = RLConfig_arguments, + RLConfig_extra_args = RLConfig_extra_args, + RLConfig_call_args = RLConfig_call_args, + + RLTrainer_extras = RLTrainer_extras, + ) + + # Create new function + created_module = create_new_function( + f"Unsloth{RLTrainer_name}", + RLTrainer_source, + f"trl.trainer.{trainer_file}", + imports, + ) + + # Patch Trainer + exec(f"trl.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}", locals(), globals()) + exec(f"trl.trainer.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}", locals(), globals()) + exec(f"trl.trainer.{trainer_file}.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}", locals(), globals()) + + # Patch Config + exec(f"trl.{RLConfig_name} = created_module.Unsloth{RLConfig_name}", locals(), globals()) + exec(f"trl.trainer.{RLConfig_name} = created_module.Unsloth{RLConfig_name}", locals(), globals()) + exec(f"trl.trainer.{trainer_file}.{RLConfig_name} = created_module.Unsloth{RLConfig_name}", locals(), globals()) +pass + + +def patch_vllm(trainer_file, RLTrainer_name, all_imports, imports): + RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}") + init = inspect.getsource(RLTrainer.__init__) + old_init = init + + # 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") + + # Set use_vllm if not set + init = re.sub( + r"\)([ ]{0,}\-\>[ ]{0,}None[ ]{0,}):\n([\s]{4})", + r"):\n\2 "\ + r"if hasattr(model, 'vllm_engine') and "\ + r"getattr(args, 'use_vllm') and getattr(args, 'use_vllm', False): "\ + r"args.use_vllm = True\n\2", + init, 1, + ) - # Replace vLLM sections since we already have it done! vllm_part = re.findall( - r"(\n[\s]{4}"\ + r"(\n[\s]{8}"\ r"if (self|args)\.use_vllm\:.+?"\ - r"\n[\s]{4,}"\ + r"\n[\s]{8,}"\ "else:\n)", - __init__, + init, flags = re.MULTILINE | re.DOTALL, ) - if (len(vllm_part) != 1): return + if len(vllm_part) != 1: return None vllm_part, args = vllm_part[0][0], vllm_part[0][1] # Strip all comments @@ -303,40 +358,31 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): new_vllm_part, flags = re.MULTILINE | re.DOTALL, ) - if len(sampling_params) != 1: return + if len(sampling_params) != 1: return None sampling_params = sampling_params[0] # Replace with our vLLM engine sampling_params = \ - " "*8 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \ + " "*12 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \ 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") - - # Add spaces back into __init__ - __init__ = __init__.split("\n") - __init__ = "\n".join(' '*spaces + x for x in __init__) + new_vllm_part = \ + f"\n{' '*8}if {args}.use_vllm:\n{sampling_params} "\ + f"if getattr(args, 'sampling_params', None) is None else "\ + f"getattr(args, 'sampling_params', None)\n{' '*8}else:\n" + init = init.replace(vllm_part, new_vllm_part) # 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__" : (old__init__, __init__,)} + changed = {"__init__" : (old_init, init,)} + for function in functions: if not hasattr(RLTrainer, function): continue fx = getattr(RLTrainer, function) - try: - source = inspect.getsource(fx) - except: - continue + try: source = inspect.getsource(fx) + except: continue original_source = source # llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model @@ -386,22 +432,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): RLTrainer_source = RLTrainer_source.replace(old, new) pass RLTrainer_source = RLTrainer_source.replace( - f"class {RLTrainer_name}", f"class Unsloth{RLTrainer_name}", 1 + f"class {RLTrainer_name}", f"class _Unsloth{RLTrainer_name}", 1 ) - - # Create new class in compiled cache and import it - module = create_new_function( - RLTrainer_name, - RLTrainer_source, - f"trl.trainer.{trainer_file}", - imports, - ) - - # Patch over modules - exec(f"trl.{RLTrainer_name} = module.Unsloth{RLTrainer_name}", locals(), globals()) - exec(f"trl.trainer.{RLTrainer_name} = module.Unsloth{RLTrainer_name}", locals(), globals()) - exec(f"trl.trainer.{trainer_file}.{RLTrainer_name} = module.Unsloth{RLTrainer_name}", locals(), globals()) - return module + return RLTrainer_source pass