From 07488965eb7c70c3d7ff1f101150c001f7cc9ab2 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Tue, 19 Mar 2024 01:37:11 +1100
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c31cb3d84b..99dc8ff22d 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -894,6 +894,16 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
 pass
 
 
+def _wrap_fast_inference(generate, device_type, dtype):
+    # Wraps inference with bfloat16 / float16
+    @torch.inference_mode
+    def _fast_generate(*args, **kwargs):
+        with torch.autocast(device_type = device_type, dtype = dtype):
+            return generate(*args, **kwargs)
+    return _fast_generate
+pass
+
+
 class FastLlamaModel:
 
     @staticmethod
@@ -1582,6 +1592,15 @@ class FastLlamaModel:
             internal_model.gradient_checkpointing = False
             internal_model.training = False
         pass
+
+        # Also check if lm_head / embeddings are trained
+        lm_head = getattr(model, "model", model).lm_head.weight
+        device_type = lm_head.device.type
+        dtype = model.config.torch_dtype
+
+        # Wrap model.generate
+        model._unwrapped_old_generate = model.generate
+        model.generate = _wrap_fast_inference(model.generate, device_type, dtype)
     pass
 
 
@@ -1602,5 +1621,14 @@ class FastLlamaModel:
             internal_model.gradient_checkpointing = use_gradient_checkpointing
             internal_model.training = True
         pass
+
+        # Also revert model.generate
+        if hasattr(model, "_unwrapped_old_generate"):
+            model.generate = model._unwrapped_old_generate
+            del model._unwrapped_old_generate
+        pass
     pass
 pass
+
+
+