From 30ddeca3e2e89aa498013944df4d3188e1653b63 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Tue, 19 Mar 2024 19:31:07 +1100
Subject: [PATCH] whoami

---
 unsloth/models/llama.py | 12 +++++++++++-
 unsloth/save.py         | 28 ++++++++++++++--------------
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 574385fba0..f34ad486d5 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -593,7 +593,17 @@ def LlamaModel_fast_forward(
     all_self_attns = () if output_attentions else None
     next_decoder_cache = () if use_cache else None
 
-    for idx, decoder_layer in enumerate(self.layers):
+    # Gradient checkpointing methods (ie sqrt)
+    if hasattr(self, "_gradient_checkpointing_boundaries"):
+        boundaries = self._gradient_checkpointing_boundaries
+    else:
+        boundaries = None
+    pass
+
+    # Find which layers to do gradient checkpointing
+    n_layers = len(self.layers)
+    
+    for idx, decoder_layer in enumerate(n_layers):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
diff --git a/unsloth/save.py b/unsloth/save.py
index 499fbb67a1..b8c2c8c29f 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -24,6 +24,7 @@ from transformers.models.llama.modeling_llama import logger
 from .kernels import fast_dequantize, QUANT_STATE, get_lora_parameters
 import subprocess
 import psutil
+import re
 
 __all__ = [
     "print_quantization_methods",
@@ -176,19 +177,6 @@ def unsloth_save_model(
     temporary_location   : str = "_unsloth_temporary_saved_buffers",
     maximum_memory_usage : float = 0.9,
 ):
-    # First check for a token!
-    if push_to_hub:
-        from huggingface_hub import whoami
-        try: 
-            username = whoami(token = token)["name"]
-        except:
-            raise RuntimeError(
-                "Unsloth: Please supply a token!\n"\
-                "Go to https://huggingface.co/settings/tokens"
-            )
-        pass
-    pass
-
     if commit_message is None: commit_message = ""
     if "Unsloth" not in commit_message:
         commit_message += " (Trained with Unsloth)"
@@ -215,7 +203,19 @@ def unsloth_save_model(
     for deletion in ("model", "tokenizer", "save_method", "temporary_location", "maximum_memory_usage"):
         del save_pretrained_settings[deletion]
     pass
-    import re
+
+    # First check for a token!
+    if push_to_hub:
+        from huggingface_hub import whoami
+        try: 
+            username = whoami(token = token)["name"]
+        except:
+            raise RuntimeError(
+                "Unsloth: Please supply a token!\n"\
+                "Go to https://huggingface.co/settings/tokens"
+            )
+        pass
+    pass
 
     assert(maximum_memory_usage > 0 and maximum_memory_usage <= 0.95)