From 2a3d4f3a8d7685f1334529a85c9d2ed1ffe97b1d Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 22 Jan 2024 04:18:42 +1100 Subject: [PATCH] fast inference --- unsloth/kernels/__init__.py | 2 +- unsloth/kernels/rope_embedding.py | 10 ++- unsloth/kernels/utils.py | 73 ++++++++++++++- unsloth/models/llama.py | 145 +++++++++++++++++++++++++----- unsloth/models/mapper.py | 6 ++ 5 files changed, 206 insertions(+), 30 deletions(-) diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index 5de19c86c5..8a3e31a32a 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -22,4 +22,4 @@ from .fast_lora import ( apply_lora_qkv, apply_lora_o, ) -from .utils import fast_dequantize, QUANT_STATE +from .utils import fast_dequantize, QUANT_STATE, fast_gemv diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index 2bf7c1b272..a9527520ab 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -134,8 +134,9 @@ class Slow_RoPE_Embedding(torch.autograd.Function): half = Q.shape[-1]//2 RH_Q = torch.cat((-Q[..., half:], Q[..., :half]), dim = -1) Q *= cos - RH_Q *= sin - Q += RH_Q + Q.addcmul_(RH_Q, sin) + # RH_Q *= sin + # Q += RH_Q ctx.save_for_backward(cos, sin) return Q pass @@ -147,8 +148,9 @@ class Slow_RoPE_Embedding(torch.autograd.Function): half = dY.shape[-1]//2 RH_dY = torch.cat((dY[..., half:], -dY[..., :half]), dim = -1) dY *= cos - RH_dY *= sin - dY += RH_dY + dY.addcmul_(RH_dY, sin) + # RH_dY *= sin + # dY += RH_dY return dY, None, None, None pass pass diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 8a7722fabd..86136a74f8 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -33,14 +33,18 @@ import bitsandbytes as bnb get_ptr = bnb.functional.get_ptr import ctypes import torch -cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 -cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4 -cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 +cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 +cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4 +cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 +cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16 +cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16 + def QUANT_STATE(W): return getattr(W, "quant_state", None) pass + def fast_dequantize(W, quant_state = None, out = None): if quant_state is None: return W if type(quant_state) is not list: @@ -90,3 +94,66 @@ def fast_dequantize(W, quant_state = None, out = None): is_transposed = (True if W.shape[0] == 1 else False) return out.t() if is_transposed else out pass + + +def fast_gemv(X, W, quant_state, out = None, out_W = None): + quant_state = W.quant_state + bsz = 1 + q_len = 1 + hd = X.shape[0] + + if type(quant_state) is not list: + # https://github.com/TimDettmers/bitsandbytes/pull/763/files + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + stats = quant_state.code + offset = quant_state.offset + state2 = quant_state.state2 + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize + else: + absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state + offset, state2 = compressed_stats + absmax2, code2, blocksize2, _, _, _, _ = state2 + pass + bout = shape[0] + if out is None: out = torch.empty(bout, dtype = dtype, device = "cuda") + else: assert(out.shape[0] == bout) + + n = 1 + m = shape[0] + k = shape[1] + lda = shape[0] + ldc = shape[0] + ldb = (X.shape[-1]+1)//2 + m = ctypes.c_int32(m) + n = ctypes.c_int32(n) + k = ctypes.c_int32(k) + lda = ctypes.c_int32(lda) + ldb = ctypes.c_int32(ldb) + ldc = ctypes.c_int32(ldc) + + df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda") + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), + ctypes.c_int(blocksize2), ctypes.c_int(df.numel()), + ) + df += offset + absmax = df + + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ + cgemm_4bit_inference_naive_bf16 + + ptr_W = get_ptr(W) + ptr_absmax = get_ptr(absmax) + ptr_stats = get_ptr(stats) + blocksize = ctypes.c_int32(blocksize) + + fx(m, n, k, get_ptr(X), ptr_W, ptr_absmax, ptr_stats, get_ptr(out), + lda, ldb, ldc, blocksize) + + return out +pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 55527dbca9..fc45368339 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -102,22 +102,46 @@ def LlamaAttention_fast_forward_inference( This means we can pass in a row of Q, but we need to remember K and V, which are called the KV cache. """ - Xn = hidden_states - bsz, _, _ = hidden_states.size() - K1, V1 = past_key_value - n_heads = self.num_heads n_groups = self.num_key_value_groups n_kv_heads = self.num_key_value_heads head_dim = self.head_dim - assert(n_kv_heads * n_groups == n_heads) + # assert(n_kv_heads * n_groups == n_heads) - Qn = self.q_proj(Xn) - Kn = self.k_proj(Xn) - Vn = self.v_proj(Xn) - Qn = Qn.view(bsz, 1, n_heads, head_dim).transpose(1, 2) - Kn = Kn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) - Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) + Xn = hidden_states.view(n_heads, head_dim) + K1, V1 = past_key_value + + # LoRA or general matrix multiplication + dtype = Xn.dtype + q_proj = self.q_proj + k_proj = self.k_proj + v_proj = self.v_proj + QW, QW_quant, QA, QB, QS = get_lora_parameters(q_proj) + KW, KW_quant, KA, KB, KS = get_lora_parameters(k_proj) + VW, VW_quant, VA, VB, VS = get_lora_parameters(v_proj) + + Qn = fast_gemv(Xn, QW, QW_quant) + Kn = fast_gemv(Xn, KW, KW_quant) + Vn = fast_gemv(Xn, VW, VW_quant) + if QA is not None: + temp_lora = torch.matmul(Xn, QA.to(dtype).t()) + Qn.addmv_(QB.to(dtype).t(), temp_lora, alpha = QS) + pass + if KA is not None: + temp_lora = torch.matmul(Xn, KA.to(dtype).t()) + Kn.addmv_(KB.to(dtype).t(), temp_lora, alpha = KS) + pass + if VA is not None: + temp_lora = torch.matmul(Xn, VA.to(dtype).t()) + Vn.addmv_(VB.to(dtype).t(), temp_lora, alpha = VS) + pass + + # Qn = self.q_proj(Xn) + # Kn = self.k_proj(Xn) + # Vn = self.v_proj(Xn) + Qn = Qn.view(1, 1, n_heads, head_dim).transpose(1, 2) + Kn = Kn.view(1, 1, n_kv_heads, head_dim).transpose(1, 2) + Vn = Vn.view(1, 1, n_kv_heads, head_dim).transpose(1, 2) kv_seq_len = K1.shape[-2] + 1 cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) @@ -130,32 +154,72 @@ def LlamaAttention_fast_forward_inference( # Grouped query attention if n_groups != 1: _, _, cached_len, _ = Kn.shape - Knn = Kn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim) - Vnn = Vn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim) - Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) - Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) + Knn = Kn[:, :, None, :, :].expand(1, n_kv_heads, n_groups, cached_len, head_dim) + Vnn = Vn[:, :, None, :, :].expand(1, n_kv_heads, n_groups, cached_len, head_dim) + Knn = Knn.reshape(1, n_heads, cached_len, head_dim) + Vnn = Vnn.reshape(1, n_heads, cached_len, head_dim) else: Knn, Vnn = Kn, Vn # Attention A = torch.matmul(Qn, Knn.transpose(2, 3)) A *= 1.0 / (self.head_dim**0.5) - A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(A.dtype) - A = torch.matmul(A, Vnn) + A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype) + A = torch.matmul(A, Vnn, out = Qn) A = A.transpose(1, 2) - A = A.reshape(bsz, 1, self.hidden_size) - A = self.o_proj(A) + A = A.reshape(1, self.hidden_size) + + # A = self.o_proj(A) + o_proj = self.o_proj + OW, OW_quant, OA, OB, OS = get_lora_parameters(o_proj) + + On = fast_gemv(A, OW, OW_quant) + if OA is not None: + temp_lora = torch.matmul(A, OA.to(dtype).t()) + On.addmv_(OB.to(dtype).t(), temp_lora, alpha = OS) + pass + A = On.reshape(1, 1, self.hidden_size) + return A, (Kn, Vn) pass torch_silu = torch.nn.functional.silu def fast_mlp_inference(self, X): - gate = self.gate_proj(X) - up = self.up_proj(X) + X = X.view(1, self.hidden_size) + dtype = X.dtype + gate_proj = self.gate_proj + up_proj = self.up_proj + down_proj = self.down_proj + + # gate = gate_proj(X) + # up = up_proj(X) + gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(gate_proj) + upW, upW_quant, upA, upB, upS = get_lora_parameters(up_proj) + downW, downW_quant, downA, downB, downS = get_lora_parameters(down_proj) + + gate = fast_gemv(X, gateW, gateW_quant) + up = fast_gemv(X, upW, upW_quant) + if gateA is not None: + temp_lora = torch.matmul(X, gateA.to(dtype).t()) + gate.addmv_(gateB.to(dtype).t(), temp_lora, alpha = gateS) + pass + if upA is not None: + temp_lora = torch.matmul(X, upA.to(dtype).t()) + up.addmv_(upB.to(dtype).t(), temp_lora, alpha = upS) + pass + gate = torch_silu(gate, inplace = True) gate *= up - X = self.down_proj(gate) + + # X = down_proj(gate) + down = fast_gemv(gate, downW, downW_quant) + if downA is not None: + temp_lora = torch.matmul(gate, downA.to(dtype).t()) + down.addmv_(downB.to(dtype).t(), temp_lora, alpha = downS) + pass + X = down.view(1, 1, self.hidden_size) + return X pass @@ -1076,4 +1140,41 @@ class FastLlamaModel: internal_model.max_seq_length = max_seq_length return model pass + + + @staticmethod + def for_inference(model): + if not hasattr(model, "_original_forward"): + model._original_forward = model.forward + pass + model.forward = torch.inference_mode(model._original_forward) + + internal_model = model + internal_model.gradient_checkpointing = False + internal_model.training = False + + while hasattr(internal_model, "model"): + internal_model = internal_model.model + internal_model.gradient_checkpointing = False + internal_model.training = False + pass + pass + + + @staticmethod + def for_training(model, use_gradient_checkpointing = True): + if hasattr(model, "_original_forward"): + model.forward = model._original_forward + pass + + internal_model = model + internal_model.gradient_checkpointing = use_gradient_checkpointing + internal_model.training = True + + while hasattr(internal_model, "model"): + internal_model = internal_model.model + internal_model.gradient_checkpointing = use_gradient_checkpointing + internal_model.training = True + pass + pass pass diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 124eaf7b27..7da9e1042e 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -42,6 +42,12 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/tinyllama", "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", ), + "unsloth/mistral-7b-instruct-v0.1-bnb-4bit" : ( + "mistralai/Mistral-7B-Instruct-v0.1", + ), + "unsloth/mistral-7b-instruct-v0.2-bnb-4bit" : ( + "mistralai/Mistral-7B-Instruct-v0.2", + ), } INT_TO_FLOAT_MAPPER = {}