From bb9539cc8218ff9e9d8c13bb790a59785f3b9fc6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 19 Aug 2024 16:17:52 -0700 Subject: [PATCH] Fix NEFTune (#937) * untrained tokens llama 3.1 base * Update tokenizer_utils.py * Update tokenizer_utils.py * Bug fixes * Update llama.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update llama.py * Update llama.py * Update llama.py --- unsloth/models/_utils.py | 28 ++++++++++++++++++++++++---- unsloth/models/llama.py | 26 ++++++++++++++++++++++++++ unsloth/tokenizer_utils.py | 14 ++++++++++++++ 3 files changed, 64 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 0c0057496b..d8904aa12b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -595,7 +595,6 @@ def _get_statistics(statistics = None, force_download = True): # You can disable this by commenting the below out try: n_cpus = psutil.cpu_count(logical = False) - keynames = "\n" + "\n".join(os.environ.keys()) if statistics is not None: pass elif "\nCOLAB_" in keynames and n_cpus == 1: statistics = "colab" @@ -604,10 +603,31 @@ def _get_statistics(statistics = None, force_download = True): elif "\nRUNPOD_" in keynames: statistics = "runpod" elif "\nAWS_" in keynames: statistics = "aws" elif "\nAZURE_" in keynames: statistics = "azure" - elif "\nK_" in keynames or "\nFUNCTION_" in keynames: statistics = "gcp" + # elif "\nK_" in keynames or "\nFUNCTION_" in keynames: statistics = "gcp" elif "\nINVOCATION_ID" in keynames: statistics = "lambda" - else: statistics = "other" - + # else: statistics = "other" + else: + def try_vllm_check(): + vendor_files = ( + "/sys/class/dmi/id/product_version", + "/sys/class/dmi/id/bios_vendor", + "/sys/class/dmi/id/product_name", + "/sys/class/dmi/id/chassis_asset_tag", + "/sys/class/dmi/id/sys_vendor", + ) + from pathlib import Path + for vendor_file in vendor_files: + path = Path(vendor_file) + if path.is_file(): + file_content = path.read_text().lower() + if "amazon" in file_content: return "aws" + elif "microsoft corporation" in file_content: return "azure" + elif "google" in file_content: return "gcp" + return "other" + pass + try: statistics = try_vllm_check() + except: statistics = "other" + pass if statistics is not None: from transformers import AutoModelForCausalLM stats_model = AutoModelForCausalLM.from_pretrained( diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6a23335c8c..048ba69193 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2234,6 +2234,9 @@ class FastLlamaModel: internal_model.gradient_checkpointing = False internal_model.training = False pass + if hasattr(internal_model, "training"): + internal_model.training = False + pass # Also check if lm_head / embeddings are trained internal_model = model @@ -2267,6 +2270,16 @@ class FastLlamaModel: internal_model._saved_temp_tokenizer.padding_side = "left" pass + # Also disable training for embeddings for NEFTune + if hasattr(model, "get_input_embeddings"): + embeddings = model.get_input_embeddings() + if hasattr(embeddings, "training"): embeddings.training = False + pass + if hasattr(model, "get_output_embeddings"): + embeddings = model.get_output_embeddings() + if hasattr(embeddings, "training"): embeddings.training = False + pass + return model pass @@ -2288,6 +2301,9 @@ class FastLlamaModel: internal_model.gradient_checkpointing = use_gradient_checkpointing internal_model.training = True pass + if hasattr(internal_model, "training"): + internal_model.training = True + pass # Also revert model.generate if hasattr(model, "_unwrapped_old_generate"): @@ -2307,6 +2323,16 @@ class FastLlamaModel: internal_model._saved_temp_tokenizer.padding_side = "right" pass + # Also re-enable training for embeddings for NEFTune + if hasattr(model, "get_input_embeddings"): + embeddings = model.get_input_embeddings() + if hasattr(embeddings, "training"): embeddings.training = True + pass + if hasattr(model, "get_output_embeddings"): + embeddings = model.get_output_embeddings() + if hasattr(embeddings, "training"): embeddings.training = True + pass + return model pass pass diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 7316656b2a..b677f864a4 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1109,6 +1109,7 @@ 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 def patch_sft_trainer_tokenizer(): """ @@ -1173,6 +1174,19 @@ def patch_sft_trainer_tokenizer(): "\n"\ "fix_untrained_tokens(self.model, self.tokenizer, self.train_dataset, eps = 1e-16)\n\n" + # Add NEFTune since it doesn't seem to work?? We need to manually inject it + check_text += \ + "\n"\ + "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"\ + "\n"\ + "if getattr(self, 'neftune_noise_alpha', None) is not None:\n"\ + " self.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"\ + "\n" + check_text = check_text.split("\n") check_text = "\n".join(" "*where + x for x in check_text)