Better TRL handling
This commit is contained in:
Daniel Han
parent
dbf46781c9
commit
ff4d021c72
1 changed files with 264 additions and 231 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue