diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 5583702e7e..6a80491922 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -15,6 +15,7 @@ import torch import gc import math +from functools import partial from typing import Optional, Tuple, List, Union from ._utils import * from ._utils import __version__ @@ -1802,8 +1803,6 @@ class FastLlamaModel: model = convert_vllm_to_huggingface(quant_state_dict, model_config, dtype) model.vllm_engine = llm model.fast_generate = model.vllm_engine.generate - - from functools import partial model.fast_generate_batches = partial(generate_batches, model.vllm_engine) pass # Return old flag @@ -2632,6 +2631,10 @@ class FastLlamaModel: gc.collect() torch.cuda.empty_cache() pass + + # Add for_inference and for_training + model.for_training = partial(FastLlamaModel.for_training, model) + model.for_inference = partial(FastLlamaModel.for_inference, model) return model pass diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 4c488187c9..ec1d65ba42 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -71,6 +71,16 @@ def PatchRL(FastLanguageModel): pass +# Handles NEFTune +def neftune_post_forward_hook(module, input, output): + if module.training: + dims = torch.tensor(output.size(1) * output.size(2)) + mag_norm = module.neftune_noise_alpha / torch.sqrt(dims) + output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm) + return output +pass + + RLTrainer_replacement = ''' import os from typing import * @@ -106,6 +116,7 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if args is None: args = Unsloth{RLConfig_name}() {RLTrainer_extra_args} super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) + {RLTrainer_post} pass ''' @@ -164,6 +175,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Process RLTrainer first arguments, call_args = processed[0] + RLTrainer_post = "" # Add tokenizer if not seen if "tokenizer" not in parameters and "processing_class" in parameters: @@ -215,7 +227,6 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): " if Version(transformers_version) <= Version('4.45.2'):\n"\ " print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\\n'\n"\ " '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`')\n" - extra_args += check_ga eval_changes = \ @@ -243,6 +254,29 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): extra_args += length_check pass + # Check NEFTune + if "neftune_noise_alpha" in call_args: + neftune_check = \ + "if hasattr(self, 'neftune_hook_handle'):\n"\ + " self.neftune_hook_handle.remove()\n"\ + " if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle\n"\ + "if getattr(args, 'neftune_noise_alpha', None) is not None:\n"\ + " model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha\n"\ + " self.neftune_hook_handle = self.model.get_input_embeddings().register_forward_hook(neftune_post_forward_hook)\n"\ + "pass\n" + RLTrainer_post += neftune_check + pass + + # Enable for training and move padding side of tokenizer to right + RLTrainer_post += \ + "if model is not None and hasattr(model, 'for_training'):\n"\ + " model.for_training()\n"\ + "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n"\ + "if 'processing_class' in locals():\n"\ + " if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n"\ + " if hasattr(processing_class, tokenizer) and hasattr(processing_class.tokenizer, 'padding_side'): "\ + "processing_class.tokenizer.padding_side = 'right'\n" + # Add statistics as well! extra_args += \ "from unsloth_zoo.logging_utils import PatchRLStatistics\n"\ @@ -251,6 +285,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Create RLTrainer args extra_args = extra_args.split("\n") extra_args = "\n".join(" "*8 + x for x in extra_args) + RLTrainer_post = RLTrainer_post.split("\n") + RLTrainer_post = "\n".join(" "*8 + x for x in RLTrainer_post) RLTrainer_arguments = arguments RLTrainer_extra_args = extra_args RLTrainer_call_args = call_args @@ -344,6 +380,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): RLConfig_kwargs = ",**kwargs"[1 if RLConfig_call_args .endswith(",") else 0:], RLTrainer_extras = RLTrainer_extras, + RLTrainer_post = RLTrainer_post, ) # Create new function diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index ab3878613f..0300d1330e 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -945,7 +945,6 @@ def patch_sft_trainer_tokenizer(): " from functools import partial\n"\ " tokenizer = partial(tokenizer, add_special_tokens = False)\n"\ " processing_class = tokenizer\n"\ - " print(1111)\n" "else:\n"\ " add_special_tokens = False if has_bos_token_already else add_special_tokens\n\n"