From e0bad0eec5ecb7aeb3916fd51ce4a6b5b5db14fa Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Tue, 30 Jan 2024 04:10:14 +1100
Subject: [PATCH] Fix inference

---
 unsloth/kernels/utils.py | 15 +++++++++------
 unsloth/models/llama.py  | 12 +++++++++---
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py
index c71dce0fff..5b31d1e5ba 100644
--- a/unsloth/kernels/utils.py
+++ b/unsloth/kernels/utils.py
@@ -191,7 +191,7 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None):
 
     W, W_quant, lora_A, lora_B, lora_S = get_lora_parameters(proj)
 
-    bsz = X.shape[0]
+    bsz, _, in_dim = X.shape
 
     if W_quant is None:
         out = torch.matmul(X, W.t())
@@ -205,15 +205,18 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None):
 
     # Add in LoRA weights
     if lora_A is not None:
+        out_dim = out.shape[2]
         dtype = X.dtype
         if bsz == 1:
-            temp_lora = torch.mv(lora_A.to(dtype), out.ravel(), out = temp_lora)
-            out.addmv_(lora_B.to(dtype).t(), temp_lora, alpha = lora_S)
+            out = out.view(out_dim)
+            temp_lora = torch.mv(lora_A.to(dtype), X.ravel(), out = temp_lora)
+            out.addmv_(lora_B.to(dtype), temp_lora, alpha = lora_S)
         else:
-            print(X.shape)
-            temp_lora = torch.matmul(out, lora_A.to(dtype).t(), out = temp_lora)
-            out.addmm_(lora_B.to(dtype).t(), temp_lora, alpha = lora_S)
+            out = out.view(bsz, out_dim)
+            temp_lora = torch.mm(X.view(bsz, in_dim), lora_A.to(dtype).t(), out = temp_lora)
+            out.addmm_(temp_lora, lora_B.to(dtype).t(), alpha = lora_S)
         pass
+        out = out.view(bsz, 1, out_dim)
     pass
 
     return out
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index dd9bf3e9e3..9de65be04b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -196,16 +196,17 @@ def fast_mlp_inference(self, X):
 
     # X = self.down_proj(gate)
     down = fast_linear_forward(self.down_proj, gate)
-    return X
+    return down
 pass
 
 
 def fast_rms_layernorm_inference(self, X):
+    old_dtype = X.dtype
     XX = X.to(torch.float32)
     variance = XX.square().mean(-1, keepdim = True)
     variance += self.variance_epsilon
     XX *= variance.rsqrt_()
-    X[:] = XX
+    X = XX.to(old_dtype) # Must preserve due to residual
     X *= self.weight
     return X
 pass
@@ -637,7 +638,12 @@ def LlamaForCausalLM_fast_forward(
     )
 
     hidden_states = outputs[0]
-    logits = self.lm_head(hidden_states)
+    if hidden_states.shape[0] == 1:
+        logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
+        logits = logits.unsqueeze(0).unsqueeze(0)
+    else:
+        logits = self.lm_head(hidden_states)
+    pass
 
     loss = None
     if labels is not None: