From 68db1c7af7138d2eb6d7cb49fc01556a92e7d15e Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 3 Feb 2024 19:44:36 +1100 Subject: [PATCH] fast inference again --- unsloth/models/llama.py | 102 ++++++++++++++++++-------------------- unsloth/models/mistral.py | 42 ++++++++-------- 2 files changed, 70 insertions(+), 74 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8cffb33ce1..d61dbf842b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -125,54 +125,51 @@ def LlamaAttention_fast_forward_inference( # Prefill phase # if not hasattr(self, "paged_attention"): if do_prefill: - # self.paged_attention = torch.empty((KV_CACHE_INCREMENT+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda") - # self.paged_attention_K = self.paged_attention[:,0] - # self.paged_attention_V = self.paged_attention[:,1] - # self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) - # self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) - # self.temp_QA = torch.empty((2, bsz, 1, hd), dtype = dtype, device = "cuda") - # self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda") - # self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda") - # self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT), dtype = dtype, device = "cuda") + self.paged_attention = torch.empty((KV_CACHE_INCREMENT+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda") + self.paged_attention_K = self.paged_attention[:,0] + self.paged_attention_V = self.paged_attention[:,1] + self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) + self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) + self.temp_QA = torch.empty((2, bsz, 1, hd), dtype = dtype, device = "cuda") + self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda") + self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda") + self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT), dtype = dtype, device = "cuda") self.scalar = 1.0 / math_sqrt(self.head_dim) - # elif kv_seq_len >= self.paged_attention.shape[0]: - # self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim)) - # self.paged_attention_K = self.paged_attention[:,0] - # self.paged_attention_V = self.paged_attention[:,1] - # self.attention.resize_((bsz, n_heads, 1, self.attention.shape[-1]+KV_CACHE_INCREMENT)) - # pass - - Qn = self.q_proj(Xn) - Kn = self.k_proj(Xn) - Vn = self.v_proj(Xn) - # Qn = fast_linear_forward(self.q_proj, Xn)#, out = self.temp_QA[0]) - # Kn = fast_linear_forward(self.k_proj, Xn)#, out = self.temp_KV[0]) - # Vn = fast_linear_forward(self.v_proj, Xn)#, out = self.temp_KV[1]) + elif kv_seq_len >= self.paged_attention.shape[0]: + self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim)) + self.paged_attention_K = self.paged_attention[:,0] + self.paged_attention_V = self.paged_attention[:,1] + self.attention.resize_((bsz, n_heads, 1, self.attention.shape[-1]+KV_CACHE_INCREMENT)) + pass + + Qn = fast_linear_forward(self.q_proj, Xn, out = self.temp_QA[0]) + Kn = fast_linear_forward(self.k_proj, Xn, out = self.temp_KV[0]) + Vn = fast_linear_forward(self.v_proj, Xn, out = self.temp_KV[1]) 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) - cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) - Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) - # cos = self.rotary_emb.cos_cached[seq_len] - # sin = self.rotary_emb.sin_cached[seq_len] - # h = head_dim // 2 + # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) + # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) + cos = self.rotary_emb.cos_cached[seq_len] + sin = self.rotary_emb.sin_cached[seq_len] + h = head_dim // 2 - # RH_Q = self.RH_Q - # RH_Q[:,:,:,:h] = Qn[:,:,:,h:]; RH_Q[:,:,:,h:] = Qn[:,:,:,:h]; torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]); - # Qn *= cos; Qn.addcmul_(RH_Q, sin); + RH_Q = self.RH_Q + RH_Q[:,:,:,:h] = Qn[:,:,:,h:]; RH_Q[:,:,:,h:] = Qn[:,:,:,:h]; torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]); + Qn *= cos; Qn.addcmul_(RH_Q, sin); - # RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda") - # RH_K[:,:,:,:h] = Kn[:,:,:,h:]; RH_K[:,:,:,h:] = Kn[:,:,:,:h]; torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]); - # Kn *= cos; Kn.addcmul_(RH_K, sin); + RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda") + RH_K[:,:,:,:h] = Kn[:,:,:,h:]; RH_K[:,:,:,h:] = Kn[:,:,:,:h]; torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]); + Kn *= cos; Kn.addcmul_(RH_K, sin); # New KV cache - Kn = torch.cat([K1, Kn], dim = 2) - Vn = torch.cat([V1, Vn], dim = 2) - # self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3) - # self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3) - # Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3) - # Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3) + # Kn = torch.cat([K1, Kn], dim = 2) + # Vn = torch.cat([V1, Vn], dim = 2) + self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3) + self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3) + Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3) + Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3) # Grouped query attention if n_groups != 1: @@ -185,32 +182,31 @@ def LlamaAttention_fast_forward_inference( Knn, Vnn = Kn, Vn # Attention - A = torch.matmul(Qn, Knn.transpose(2, 3))#, out = self.attention[:,:,:,:kv_seq_len]) + A = torch.matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:kv_seq_len]) A *= self.scalar 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 = fast_linear_forward(self.o_proj, A)#, out = self.temp_QA[1]) + A = fast_linear_forward(self.o_proj, A, out = self.temp_QA[1]) return A, (Kn, Vn) pass def fast_mlp_inference(self, X): + # gate = self.gate_proj(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") - gate = self.gate_proj(X) - up = self.up_proj(X) - # 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 - down = self.down_proj(gate) - # down = fast_linear_forward(self.down_proj, gate)#, out = up[:,:,:hd]) + # X = self.down_proj(gate) + down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:hd]) return down pass @@ -368,7 +364,7 @@ def LlamaDecoderLayer_fast_forward( # Self Attention residual = hidden_states - hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states) + hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states) hidden_states, present_key_value = LlamaAttention_fast_forward_inference( self.self_attn, hidden_states, @@ -380,7 +376,7 @@ def LlamaDecoderLayer_fast_forward( # Fully Connected residual = hidden_states - hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) + hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) hidden_states = fast_mlp_inference(self.mlp, hidden_states) hidden_states += residual else: @@ -624,7 +620,7 @@ def LlamaModel_fast_forward_inference( for idx, decoder_layer in enumerate(self.layers): # Self Attention residual = hidden_states - hidden_states = fast_rms_layernorm(decoder_layer.input_layernorm, hidden_states) + hidden_states = fast_rms_layernorm_inference(decoder_layer.input_layernorm, hidden_states) hidden_states, present_key_value = LlamaAttention_fast_forward_inference( decoder_layer.self_attn, hidden_states, @@ -635,13 +631,13 @@ def LlamaModel_fast_forward_inference( # Fully Connected residual = hidden_states - hidden_states = fast_rms_layernorm(decoder_layer.post_attention_layernorm, hidden_states) + hidden_states = fast_rms_layernorm_inference(decoder_layer.post_attention_layernorm, hidden_states) hidden_states = fast_mlp_inference(decoder_layer.mlp, hidden_states) hidden_states += residual next_decoder_cache.append(present_key_value) pass - hidden_states = fast_rms_layernorm(self.norm, hidden_states) + hidden_states = fast_rms_layernorm_inference(self.norm, hidden_states) return BaseModelOutputWithPast( last_hidden_state = hidden_states, diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index c0bdae5d3b..cd92d0e7dd 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -201,31 +201,31 @@ def MistralForCausalLM_fast_forward( # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) self.model._has_no_labels = labels is None - # if past_key_values is not None and \ - # hasattr(self.model.layers[0].self_attn, "paged_attention"): - # outputs = LlamaModel_fast_forward_inference( - # self.model, - # input_ids, - # past_key_values, - # ) - # else: - outputs = self.model( - input_ids=input_ids, - causal_mask=causal_mask, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) + if past_key_values is not None and \ + hasattr(self.model.layers[0].self_attn, "paged_attention"): + outputs = LlamaModel_fast_forward_inference( + self.model, + input_ids, + past_key_values, + ) + else: + outputs = self.model( + input_ids=input_ids, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) pass hidden_states = outputs[0] bsz, q_len, hd = hidden_states.shape - if False:#bsz == 1 and q_len == 1: + if bsz == 1 and q_len == 1: logits = torch.mv(self.lm_head.weight, hidden_states.ravel()) logits = logits.unsqueeze(0).unsqueeze(0) else: