From e4ac52fe85ce63e86aaf158a6611a3061c219bb0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 5 Feb 2025 02:30:51 -0800
Subject: [PATCH] Update llama.py

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

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index b350f764c4..1a700c62da 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1784,7 +1784,6 @@ class FastLlamaModel:
                 gpu_memory_utilization = gpu_memory_utilization,
                 max_seq_length         = max_seq_length,
                 dtype                  = dtype,
-                disable_log_stats      = disable_log_stats,
                 float8_kv_cache        = float8_kv_cache,
                 enable_lora            = True,
                 max_lora_rank          = max_lora_rank,
@@ -2302,6 +2301,20 @@ class FastLlamaModel:
             modules_to_save = list(set(modules_to_save))
         pass
 
+        vllm_engine = None
+        if hasattr(model, "vllm_engine"):
+            # Fast inference!
+            vllm_engine = model.vllm_engine
+            vllm_fast_generate = model.fast_generate
+            vllm_fast_generate_batches = model.fast_generate_batches
+
+            if len(modules_to_save) != 0:
+                raise NotImplementedError("Unsloth: Currently fast inference does not work with training embeddings or lm_head.")
+
+            if bias != "none":
+                raise NotImplementedError("Unsloth: Currently fast inference does not work with using biases for LoRA.")
+        pass
+
         # Get LoRA
         arguments = dict(
             r                   = r,
@@ -2408,6 +2421,19 @@ class FastLlamaModel:
             torch.cuda.empty_cache()
         pass
 
+        # Patch for fast inference
+        if vllm_engine is not None:
+            model.vllm_engine = vllm_engine
+            model.fast_generate = vllm_fast_generate
+            model.fast_generate_batches = vllm_fast_generate_batches
+
+            # Also saving and loading LoRA
+            from functools import partial
+            from unsloth_zoo.vllm_utils import save_lora, load_lora
+            model.save_lora = partial(save_lora, model)
+            model.load_lora = partial(load_lora, model)
+        pass
+
         return model
     pass