fast inference
This commit is contained in:
Daniel Han-Chen
parent
a5ab4dc21a
commit
2a3d4f3a8d
5 changed files with 206 additions and 30 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue