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

@@ -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(

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

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)