From 063a2b42b02f9991f723b2c3f8637378cd3bc28d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Sep 2024 12:18:43 -0700
Subject: [PATCH] Llama 3.2

---
 unsloth/models/llama.py  | 10 ++++++++--
 unsloth/models/loader.py | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index a245330108..aedc935d75 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -62,7 +62,10 @@ from peft.tuners.lora import Linear4bit as Peft_Linear4bit
 from ..save import patch_saving_functions
 import re, os, inspect, math, sys
 from huggingface_hub.utils._token import get_token
-
+from transformers import __version__ as transformers_version
+from packaging.version import Version
+transformers_version = Version(transformers_version)
+SUPPORTS_LLAMA32 = transformers_version >= Version("4.46")
 
 def original_apply_qkv(self, X):
     Q = self.q_proj(X)
@@ -1383,7 +1386,10 @@ def _wrap_fast_inference(generate, device_type, dtype, model):
         pass
 
         # For newer HF
-        kwargs["cache_implementation"] = "dynamic"
+        if SUPPORTS_LLAMA32:
+            kwargs["cache_implementation"] = "hybrid"
+        else:
+            kwargs["cache_implementation"] = "dynamic"
         # For num_logits_to_keep
         kwargs["num_logits_to_keep"] = 1
 
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 13710eeda1..8afb3a3a82 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -31,6 +31,7 @@ SUPPORTS_FOURBIT = transformers_version >= Version("4.37")
 SUPPORTS_GEMMA   = transformers_version >= Version("4.38")
 SUPPORTS_GEMMA2  = transformers_version >= Version("4.42")
 SUPPORTS_LLAMA31 = transformers_version >= Version("4.43.2")
+SUPPORTS_LLAMA32 = transformers_version >= Version("4.46")
 if SUPPORTS_GEMMA:
     from .gemma  import FastGemmaModel
 if SUPPORTS_GEMMA2:
@@ -137,6 +138,17 @@ def get_model_name(model_name, load_in_4bit = True):
 pass
 
 
+LLAMA32_MODEL_NAMES = set((
+    "unsloth/Llama-3.2-1B-bnb-4bit",
+    "unsloth/Llama-3.2-3B-bnb-4bit",
+    "unsloth/Llama-3.2-1B",
+    "unsloth/Llama-3.2-3B",
+    "unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
+    "unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
+    "unsloth/Llama-3.2-1B-Instruct",
+    "unsloth/Llama-3.2-3B-Instruct",
+))
+
 class FastLanguageModel(FastLlamaModel):
     @staticmethod
     def from_pretrained(
@@ -242,7 +254,17 @@ class FastLanguageModel(FastLlamaModel):
                     f'Try `pip install --upgrade "transformers>=4.43.2"`\n'\
                     f"to obtain the latest transformers build, then restart this session."\
                 )
+
+            elif model_name.lower() in LLAMA32_MODEL_NAMES and not SUPPORTS_LLAMA32:
+                raise ImportError(
+                    f"Unsloth: Your transformers version of {transformers_version} does not support Llama 3.2.\n"\
+                    f"The minimum required version is 4.46\n"\
+                    f'Try `pip install --upgrade "transformers>=4.46"`\n'\
+                    f"to obtain the latest transformers build, then restart this session."\
+                )
+
             dispatch_model = FastLlamaModel
+            
         elif model_type == "mistral": dispatch_model = FastMistralModel
         elif model_type == "gemma":
             if not SUPPORTS_GEMMA: