Bug fixes

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

@@ -1628,7 +1628,7 @@ class FastLlamaModel:
 
         # Torch.compile fails on embedding matrix??
         # Workaround randomnly fixes it for torch versions < 2.
-        model.set_input_embeddings(torch.nn.Embedding.from_pretrained(model.get_input_embeddings().weight))
+        # model.set_input_embeddings(torch.nn.Embedding.from_pretrained(model.get_input_embeddings().weight))
         model.config.update({"unsloth_version" : __version__})
 
         # We also do this for the lm_head
@@ -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