This commit is contained in:
Daniel Han-Chen 2024-02-01 03:04:46 +11:00
commit e2f72fe52f
2 changed files with 8 additions and 10 deletions

View file

@ -232,9 +232,9 @@ def LlamaAttention_fast_forward_inference(
head_dim = self.head_dim
assert(n_kv_heads * n_groups == n_heads)
Qn = fast_linear_forward(self.q_proj, Xn)
Kn = fast_linear_forward(self.k_proj, Xn)
Vn = fast_linear_forward(self.v_proj, Xn)
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)
@ -260,22 +260,21 @@ def LlamaAttention_fast_forward_inference(
# Attention
A = torch.matmul(Qn, Knn.transpose(2, 3))
A *= 1.0 / math_sqrt(self.head_dim)
A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
A = torch.matmul(A, Vnn, out = Qn)
A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(A.dtype)
A = torch.matmul(A, Vnn)
A = A.transpose(1, 2)
A = A.reshape(bsz, 1, self.hidden_size)
A = fast_linear_forward(self.o_proj, A)
A = self.o_proj(A)
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)
gate = fast_linear_forward(self.gate_proj, X)
up = fast_linear_forward(self. up_proj, X)
gate = torch_silu(gate, inplace = True)
gate = torch.nn.functional.silu(gate, inplace = True)
gate *= up
# X = self.down_proj(gate)
@ -725,7 +724,7 @@ def LlamaForCausalLM_fast_forward(
hidden_states = outputs[0]
bsz, q_len, hd = hidden_states.shape
if bsz == 1 and q_len == 1:
if False:#bsz == 1 and q_len == 1:
logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
logits = logits.unsqueeze(0).unsqueeze(0)
else:

View file

@ -68,7 +68,6 @@ __INT_TO_FLOAT_MAPPER = \
"01-ai/Yi-6B",
),
"unsloth/solar-10.7b-bnb-4bit" : (
"unsloth/solar-10.7b",
"upstage/SOLAR-10.7B-v1.0",
),
}