From a81d19383047d3bfaae4e79efe171f22dd83439d Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 3 Feb 2024 02:20:00 +1100 Subject: [PATCH] Update utils.py --- unsloth/kernels/utils.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 0d21b19fdf..775d881eea 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -203,16 +203,23 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None): # Add in LoRA weights if lora_A is not None: + if not hasattr(lora_A, "_fast_lora_A"): + lora_A._fast_lora_A = lora_A.to(dtype) + lora_A._fast_lora_B = lora_B.to(dtype) + pass + lora_A = lora_A._fast_lora_A + lora_B = lora_A._fast_lora_B + out_dim = out.shape[2] dtype = X.dtype if bsz == 1: 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) + temp_lora = torch.mv(lora_A, X.ravel(), out = temp_lora) + out.addmv_(lora_B, temp_lora, alpha = lora_S) else: 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) + temp_lora = torch.mm(X.view(bsz, in_dim), lora_A.t(), out = temp_lora) + out.addmm_(temp_lora, lora_B.t(), alpha = lora_S) pass out = out.view(bsz, 1, out_dim) pass