From a5ee70b63ff70ffd7b99c49906adb4a8369cd4d5 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Thu, 1 Feb 2024 18:34:43 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 4989f30db5..3e15b82876 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -284,13 +284,17 @@ pass def fast_mlp_inference(self, X): # gate = self.gate_proj(X) # up = self.up_proj(X) - gate = fast_linear_forward(self.gate_proj, X) - up = fast_linear_forward(self. up_proj, X) + bsz, _, hd = X.shape + mlp_size = model.config.intermediate_size + # temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda") + + gate = fast_linear_forward(self.gate_proj, X)#, out = temp[0]) + up = fast_linear_forward(self. up_proj, X)#, out = temp[1]) gate = torch.nn.functional.silu(gate, inplace = True) gate *= up # X = self.down_proj(gate) - down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:X.shape[2]]) + down = fast_linear_forward(self.down_proj, gate)#, out = up[:,:,:hd]) return down pass