diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 6bd8a6f345..5135766d8b 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -45,10 +45,9 @@ def fast_geglu_inference(self, 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 = 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) + up = fast_linear_forward(self. up_proj, X) gate = torch_nn_functional_gelu(gate, approximate = "tanh") gate *= up @@ -83,20 +82,30 @@ def GemmaDecoderLayer_fast_forward( padding_mask: Optional[torch.LongTensor] = None, *args, **kwargs, ): - if past_key_value is not None: + if use_cache: #past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention") out_weight = torch.empty(self.input_layernorm.weight.shape, dtype = torch.float32, device = "cuda") # Self Attention residual = hidden_states hidden_states = fast_rms_layernorm_inference_gemma(self.input_layernorm, hidden_states, out_weight) - hidden_states, present_key_value = LlamaAttention_fast_forward_inference( - self.self_attn, - hidden_states, - past_key_value, - position_ids, - do_prefill = do_prefill, + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + padding_mask=padding_mask, ) + # hidden_states, present_key_value = LlamaAttention_fast_forward_inference( + # self.self_attn, + # hidden_states, + # past_key_value, + # position_ids, + # do_prefill = do_prefill, + # ) hidden_states += residual # Fully Connected @@ -129,13 +138,8 @@ def GemmaDecoderLayer_fast_forward( pass outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights,) - - if use_cache: - outputs += (present_key_value,) - + if output_attentions: outputs += (self_attn_weights,) + if use_cache: outputs += (present_key_value,) return outputs pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index bfbd10eb89..b802403abc 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -74,7 +74,7 @@ pass from math import sqrt as math_sqrt -KV_CACHE_INCREMENT = 128 # KV Cache update size +KV_CACHE_INCREMENT = 256 # KV Cache update size def LlamaAttention_fast_forward_inference( self, @@ -82,6 +82,7 @@ def LlamaAttention_fast_forward_inference( past_key_value: Optional[Tuple[torch.Tensor]], position_ids, do_prefill = False, + attention_mask = None, ): """ https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406 @@ -167,12 +168,12 @@ def LlamaAttention_fast_forward_inference( 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) # Handle sliding windows sliding_window = getattr(self.config, "sliding_window", None) @@ -200,6 +201,7 @@ def LlamaAttention_fast_forward_inference( # Attention A = torch.matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len]) A *= self.scalar + if attention_mask is not None: A += attention_mask # Must add attention_mask for batched 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) @@ -215,10 +217,9 @@ def fast_swiglu_inference(self, 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 = 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) + up = fast_linear_forward(self. up_proj, X) gate = torch_nn_functional_silu(gate, inplace = True) gate *= up @@ -375,19 +376,30 @@ def LlamaDecoderLayer_fast_forward( (see `past_key_values`). past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states """ - if past_key_value is not None: + if use_cache: #past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention") # Self Attention residual = 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, - past_key_value, - position_ids, - do_prefill = do_prefill, + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + padding_mask=padding_mask, ) + # hidden_states, present_key_value = LlamaAttention_fast_forward_inference( + # self.self_attn, + # hidden_states, + # past_key_value, + # position_ids, + # do_prefill = do_prefill, + # attention_mask = attention_mask, + # ) hidden_states += residual # Fully Connected @@ -418,13 +430,8 @@ def LlamaDecoderLayer_fast_forward( pass outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights,) - - if use_cache: - outputs += (present_key_value,) - + if output_attentions: outputs += (self_attn_weights,) + if use_cache: outputs += (present_key_value,) return outputs pass @@ -602,9 +609,8 @@ def LlamaModel_fast_forward( pass for idx, decoder_layer in enumerate(self.layers): - if output_hidden_states: - all_hidden_states += (hidden_states,) + if output_hidden_states: all_hidden_states += (hidden_states,) past_key_value = past_key_values[idx] if past_key_values is not None else None if self.gradient_checkpointing and self.training: @@ -636,22 +642,15 @@ def LlamaModel_fast_forward( use_cache=use_cache, padding_mask=padding_mask, ) + pass hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + if use_cache: next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) + if output_attentions: all_self_attns += (layer_outputs[1],) pass - hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - + if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None if not return_dict: return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) @@ -670,12 +669,29 @@ def LlamaModel_fast_forward_inference( self, input_ids, past_key_values, + attention_mask = None, ): # Fix out of bounds tokenization input_ids = input_ids[:,:self.max_seq_length] hidden_states = self.embed_tokens(input_ids) hidden_states = hidden_states.to(self.config.torch_dtype) + bsz, q_len, hd = hidden_states.shape + seq_len = past_key_values[0][0].shape[-2] + + # Must use attention mask for batched processing + sliding_window = getattr(self.config, "sliding_window", None) + if (sliding_window is not None and seq_len >= sliding_window) or (bsz != 1): + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (bsz, q_len), + hidden_states, + seq_len, + sliding_window = sliding_window, + ) + else: + attention_mask = None + pass next_decoder_cache = [] for idx, decoder_layer in enumerate(self.layers): @@ -686,7 +702,9 @@ def LlamaModel_fast_forward_inference( decoder_layer.self_attn, hidden_states, past_key_values[idx], - None, + position_ids = None, + do_prefill = False, + attention_mask = attention_mask, ) hidden_states += residual @@ -726,11 +744,12 @@ def CausalLM_fast_forward(fast_forward_inference): *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: - if past_key_values is not None and hasattr(self.model.layers[0].self_attn, "paged_attention"): + if False:#past_key_values is not None and hasattr(self.model.layers[0].self_attn, "paged_attention"): outputs = fast_forward_inference( self.model, input_ids, past_key_values, + attention_mask = attention_mask, ) else: causal_mask = xformers.attn_bias.LowerTriangularMask() diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index c609d2ecad..fcc1ab62e0 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -200,12 +200,13 @@ 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 \ + if False and 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, + attention_mask = attention_mask, ) else: outputs = self.model(