From e8ec80a4c25832ebe88c70850adc2fef310d3abf Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Thu, 1 Feb 2024 19:19:44 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 53dec892fb..d7a4499c1a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -286,15 +286,15 @@ def fast_mlp_inference(self, X): # up = self.up_proj(X) bsz, _, hd = X.shape mlp_size = self.config.intermediate_size - # temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda") + 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 = 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[:,:,:hd]) + down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:hd]) return down pass