From cbc1c69e294bf6a491f77bde13b30ac2fa9dd454 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Sun, 21 Jan 2024 16:22:28 +1100
Subject: [PATCH] faster saving & inference

---
 unsloth/models/llama.py |  1 +
 unsloth/save.py         | 64 ++++++++++++++++++++++++++++++++---------
 2 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index bc7aeb3e59..9058bc385b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -327,6 +327,7 @@ def LlamaDecoderLayer_fast_forward(
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
     else:
+        print(hidden_states.shape)
         # Self Attention
         residual = hidden_states
         hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
diff --git a/unsloth/save.py b/unsloth/save.py
index a3d5fe9b0c..0a48866b3e 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -783,12 +783,27 @@ def unsloth_save_pretrained_gguf(
     del arguments["quantization_method"]
 
     # Non blocking install GGUF first
-    git_clone = install_llama_cpp_clone_non_blocking()
-    python_install = install_python_non_blocking(["gguf", "protobuf"])
-    git_clone.wait()
-    makefile  = install_llama_cpp_make_non_blocking()
-    new_save_directory = unsloth_save_model(**arguments)
-    python_install.wait()
+    if not os.path.exists("llama.cpp"):
+        git_clone = install_llama_cpp_clone_non_blocking()
+        python_install = install_python_non_blocking(["gguf", "protobuf"])
+        git_clone.wait()
+        makefile  = install_llama_cpp_make_non_blocking()
+        new_save_directory = unsloth_save_model(**arguments)
+        python_install.wait()
+    else:
+        try:
+            new_save_directory = unsloth_save_model(**arguments)
+            makefile = None
+        except:
+            # Retry by recloning llama.cpp
+            git_clone = install_llama_cpp_clone_non_blocking()
+            python_install = install_python_non_blocking(["gguf", "protobuf"])
+            git_clone.wait()
+            makefile  = install_llama_cpp_make_non_blocking()
+            new_save_directory = unsloth_save_model(**arguments)
+            python_install.wait()
+        pass
+    pass
 
     for _ in range(3):
         gc.collect()
@@ -801,7 +816,10 @@ def unsloth_save_pretrained_gguf(
             self, save_directory, token,
             "GGUF converted", "gguf", file_location,
         )
-        print(f"Saved to https://huggingface.co/{username}/{new_save_directory.lstrip('/.')}")
+        link = f"{username}/{new_save_directory.lstrip('/.')}" \
+            if username not in new_save_directory else \
+            new_save_directory.lstrip('/.')
+        print(f"Saved to https://huggingface.co/{link}")
     pass
 pass
 
@@ -863,16 +881,31 @@ def unsloth_push_to_hub_gguf(
     del arguments["quantization_method"]
 
     # Non blocking install GGUF first
-    git_clone = install_llama_cpp_clone_non_blocking()
-    python_install = install_python_non_blocking(["gguf", "protobuf"])
-    git_clone.wait()
-    makefile  = install_llama_cpp_make_non_blocking()
-    new_save_directory = unsloth_save_model(**arguments)
+    if not os.path.exists("llama.cpp"):
+        git_clone = install_llama_cpp_clone_non_blocking()
+        python_install = install_python_non_blocking(["gguf", "protobuf"])
+        git_clone.wait()
+        makefile  = install_llama_cpp_make_non_blocking()
+        new_save_directory = unsloth_save_model(**arguments)
+        python_install.wait()
+    else:
+        try:
+            new_save_directory = unsloth_save_model(**arguments)
+            makefile = None
+        except:
+            # Retry by recloning llama.cpp
+            git_clone = install_llama_cpp_clone_non_blocking()
+            python_install = install_python_non_blocking(["gguf", "protobuf"])
+            git_clone.wait()
+            makefile  = install_llama_cpp_make_non_blocking()
+            new_save_directory = unsloth_save_model(**arguments)
+            python_install.wait()
+        pass
+    pass
 
     for _ in range(3):
         gc.collect()
 
-    python_install.wait()
     file_location = save_to_gguf(new_save_directory, quantization_method, makefile)
 
     print("Unsloth: Uploading GGUF to Huggingface Hub...")
@@ -880,7 +913,10 @@ def unsloth_push_to_hub_gguf(
         self, repo_id, token,
         "GGUF converted", "gguf", file_location,
     )
-    print(f"Saved to https://huggingface.co/{username}/{new_save_directory.lstrip('/')}")
+    link = f"{username}/{new_save_directory.lstrip('/.')}" \
+        if username not in new_save_directory else \
+        new_save_directory.lstrip('/.')
+    print(f"Saved to https://huggingface.co/{link}")
 pass