diff --git a/pyproject.toml b/pyproject.toml index 59ef1b8aab..1e8fd7e236 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,8 +42,8 @@ huggingface = [ "psutil", "wheel>=0.42.0", "numpy", - "accelerate>=0.26.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3", + "accelerate>=0.34.1", + "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,<=0.11.1", "peft>=0.7.1,!=0.11.0", "protobuf<4.0.0", "huggingface_hub", @@ -224,8 +224,8 @@ colab-new = [ "hf_transfer", ] colab-no-deps = [ - "accelerate>=0.26.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3", + "accelerate>=0.34.1", + "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,<=0.11.1", "peft>=0.7.1", "xformers<0.0.27", "bitsandbytes>=0.43.3", diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 04690d3566..cdce372b50 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1135,7 +1135,39 @@ from inspect import getsource import trl.trainer.sft_trainer from trl.trainer.sft_trainer import * from transformers.trainer import * -from trl.trainer.sft_trainer import neftune_post_forward_hook +try: + from trl.trainer.sft_trainer import neftune_post_forward_hook +except: + def neftune_post_forward_hook(module, input, output): + """ + Implements the NEFTune forward pass for the model using forward hooks. Note this works only for + torch.nn.Embedding layers. This method is slightly adapted from the original source code + that can be found here: https://github.com/neelsjain/NEFTune + + Simply add it to your model as follows: + ```python + model = ... + model.embed_tokens.neftune_noise_alpha = 0.1 + model.embed_tokens.register_forward_hook(neftune_post_forward_hook) + ``` + + Args: + module (`torch.nn.Module`): + The embedding module where the hook is attached. Note that you need to set + `module.neftune_noise_alpha` to the desired noise alpha value. + input (`torch.Tensor`): + The input tensor to the model. + output (`torch.Tensor`): + The output tensor of the model (i.e. the embeddings). + """ + 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 +pass + def patch_sft_trainer_tokenizer(): """