diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 4e7fcfa7a4..3e1b6993f8 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -23,7 +23,9 @@ import os import re from unsloth_zoo.compiler import create_new_function from unsloth_zoo.logging_utils import PatchRLStatistics - +from .rl_replacements import ( + RL_EXTRA_ARGS, +) def PatchRL(FastLanguageModel): @@ -282,6 +284,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "from unsloth_zoo.logging_utils import PatchRLStatistics\n"\ f"PatchRLStatistics('{trainer_file}')\n" + # Patch optional args + if trainer_file in RL_EXTRA_ARGS: + process_extra_args = RL_EXTRA_ARGS[trainer_file] + for process_extra_arg in process_extra_args: + extra_args += process_extra_args(call_args, extra_args) + pass + # Create RLTrainer args extra_args = extra_args.split("\n") extra_args = "\n".join(" "*8 + x for x in extra_args) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py new file mode 100644 index 0000000000..56ad57f5cc --- /dev/null +++ b/unsloth/models/rl_replacements.py @@ -0,0 +1,50 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__all__ = [ + "RL_EXTRA_ARGS", +] + +RL_EXTRA_ARGS = dict() + +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"\ + "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"\ + "fix_untrained_tokens(model, tokenizer, train_dataset, IGNORED_TOKENIZER_NAMES, eps = 1e-16)\n"\ + "fix_zero_training_loss(model, tokenizer, train_dataset)\n" + return fix_tokenizer + return "" +pass +RL_EXTRA_ARGS["sft_trainer"] = [sft_trainer_fix_untraiend_tokens,] + + +def dpo_trainer_fix_columns(call_args, extra_args): + if "model" in call_args and "train_dataset" in call_args: + fix_dpo = \ + "if hasattr(train_dataset, 'column_names'):\n"\ + " column_names = set(train_dataset.column_names)\n"\ + " check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',\n"\ + " 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',\n"\ + " 'prompt_input_ids', 'prompt_attention_mask']\n"\ + " if all(x in column_names for x in check):\n"\ + " train_dataset = train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])\n"\ + " del check, column_names\n"\ + return fix_dpo + return "" +pass +RL_EXTRA_ARGS["dpo_trainer"] = [dpo_trainer_fix_columns,] diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 0300d1330e..404fce319f 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -59,6 +59,7 @@ IGNORED_TOKENIZER_NAMES = frozenset( [x.lower() for x in IGNORED_TOKENIZER_NAMES] + \ [x.lower()+"-bnb-4bit" for x in IGNORED_TOKENIZER_NAMES] ) +os.environ["UNSLOTH_IGNORED_TOKENIZER_NAMES"] = "\n".join(IGNORED_TOKENIZER_NAMES) # Check environments keynames = "\n" + "\n".join(os.environ.keys()) @@ -1055,5 +1056,5 @@ def patch_sft_trainer_tokenizer(): pass pass -# Finally patch TRL tokenizer things +# Finally patch TRL tokenizer things -> moved to RL # patch_sft_trainer_tokenizer()