diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 48c9fe803b..664a647f47 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -57,173 +57,6 @@ def fast_geglu_inference(self, X): pass -class FastGemmaRotaryEmbedding(torch.nn.Module): - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): - super().__init__() - - self.dim = dim - self.max_position_embeddings = max_position_embeddings - self.base = base - self.register_buffer("inv_freq", None, persistent=False) - self.register_buffer("cos_cached", None, persistent=False) - self.register_buffer("sin_cached", None, persistent=False) - - self.inv_freq = 1.0 / ( - self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cuda").float() / self.dim) - ) - - position_ids = torch.arange(self.max_position_embeddings, device="cuda", dtype=torch.int64).unsqueeze(0) - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1) - position_ids_expanded = position_ids[:, None, :].float() - freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - self.cos_cached2 = freqs - self.sin_cached2 = emb.sin().to(torch.bfloat16) - - def forward(self, x, position_ids, seq_len=None): - # x: [bs, num_attention_heads, seq_len, head_size] - if self.inv_freq is None: - self.inv_freq = 1.0 / ( - self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim) - ) - - # length = position_ids.shape[1] - if self.cos_cached is None: - position_ids = torch.arange(self.max_position_embeddings, device=x.device, dtype=torch.int64).unsqueeze(0) - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1) - position_ids_expanded = position_ids[:, None, :].float() - freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - self.cos_cached = emb.cos().to(dtype=x.dtype) - self.sin_cached = emb.sin().to(dtype=x.dtype) - pass - - position_ids = torch.arange(self.max_position_embeddings, device=x.device, dtype=torch.int64).unsqueeze(0) - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1) - position_ids_expanded = position_ids[:, None, :].float() - freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos_cached = emb.cos().to(dtype=x.dtype) - sin_cached = emb.sin().to(dtype=x.dtype) - - print(freqs) - print(self.cos_cached2) - - # return self.cos_cached[:,:length], self.sin_cached[:,:length] - return self.cos_cached, self.sin_cached -pass - - -# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L320 -def GemmaAttention_fast_forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_value = None, #Optional[Cache] = None, - output_attentions: bool = False, - use_cache: bool = False, - cache_position: Optional[torch.LongTensor] = None, - **kwargs, -): - # Clear inference - if hasattr(self, "paged_attention"): - del self.paged_attention_K - del self.paged_attention_V - del self.paged_attention - del self.temp_QA - del self.temp_KV - del self.RH_Q - del self.attention - pass - - bsz, q_len, _ = hidden_states.size() - - 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) - - Q, K, V = self.apply_qkv(self, hidden_states) - Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2) - K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2) - V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2) - - if True:# position_ids is None: - # cos = self.rotary_emb.cos_cached - # sin = self.rotary_emb.sin_cached - cos, sin = self.rotary_emb(V, position_ids, seq_len = q_len) - Q, K = fast_rope_embedding(Q, K, cos, sin) - else: - cos, sin = self.rotary_emb(V, position_ids, seq_len = q_len) - Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids) - pass - - past_key_value = getattr(self, "past_key_value", past_key_value) - - if past_key_value is not None: - # sin and cos are specific to RoPE models; position_ids needed for the static cache - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - K, V = past_key_value.update(K, V, self.layer_idx, cache_kwargs) - - # Attention module - if (not HAS_FLASH_ATTENTION):# and attention_mask is None): - # Xformers memory efficient attention - # Also has Flash Attention v2 dispatching - Q = Q.transpose(1, 2) - K = K.transpose(1, 2) - V = V.transpose(1, 2) - - # Group query attention - if n_groups != 1: - K = K .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim) - V = V .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim) - K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim) - V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim) - if hidden_states.requires_grad: - K = K.reshape(bsz, kv_seq_len, n_heads, head_dim) - V = V.reshape(bsz, kv_seq_len, n_heads, head_dim) - else: - Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim) - pass - A = xformers_attention(Q, K, V, attn_bias = causal_mask) - A = A.view(bsz, q_len, n_heads, head_dim) - - elif HAS_FLASH_ATTENTION and attention_mask is None: - Q = Q.transpose(1, 2) - K = K.transpose(1, 2) - V = V.transpose(1, 2) - A = flash_attn_func(Q, K, V, causal = True) - else: - causal_mask = attention_mask - if attention_mask is not None and cache_position is not None: - causal_mask = causal_mask[:, :, cache_position, : K.shape[-2]] - - # Grouped query attention - if n_groups != 1: - K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim) - V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim) - K = K.reshape(bsz, n_heads, kv_seq_len, head_dim) - V = V.reshape(bsz, n_heads, kv_seq_len, head_dim) - pass - # Must be contiguous or else results are False! - # https://github.com/pytorch/pytorch/issues/112577 - Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous() - # Needs (batch_size, n_heads, seq_len, head_dim) - # is_casual and attention_mask must not be both set! - A = scaled_dot_product_attention(Q, K, V, attn_mask = causal_mask, is_causal = False) - # Go back to (batch_size, seq_len, n_heads, head_dim) - A = A.transpose(1, 2).contiguous() - pass - - attn_output = A.reshape(bsz, q_len, n_heads*head_dim) - attn_output = self.apply_o(self, attn_output) - attn_weights = None - return attn_output, attn_weights, past_key_value -pass - - # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590 def GemmaDecoderLayer_fast_forward( self, @@ -340,227 +173,6 @@ def GemmaModel_fast_forward_inference( pass -# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825 -def GemmaModel_fast_forward( - self, - input_ids: torch.LongTensor, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - *args, **kwargs, -) -> Union[Tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError( - "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) - - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - past_seen_tokens = 0 - if use_cache: # kept for BC (cache positions) - if not isinstance(past_key_values, StaticCache): - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - past_seen_tokens = past_key_values.get_seq_length() - - if cache_position is None: - cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device - ) - - if position_ids is None: - position_ids = cache_position.unsqueeze(0) - - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds) - - # embed positions - hidden_states = inputs_embeds - - # normalized - hidden_states = hidden_states * (self.config.hidden_size**0.5) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - ) - - 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],) - - - # hidden_states = self.norm(hidden_states) - hidden_states = fast_rms_layernorm(self.norm, hidden_states) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = None - if use_cache: - next_cache = ( - next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache - ) - 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) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) -pass - - -def GemmaForCausalLM_fast_forward( - self, - input_ids: torch.LongTensor = None, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - *args, **kwargs, -) -> Union[Tuple, CausalLMOutputWithPast]: - - if causal_mask is None and past_key_values is None: - causal_mask = xformers.attn_bias.LowerTriangularMask() - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # 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 = GemmaModel_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 bsz == 1 and q_len == 1: - logits = torch.mv(self.lm_head.weight, hidden_states.ravel()) - logits = logits.unsqueeze(0).unsqueeze(0) - else: - logits = self.lm_head(hidden_states) - pass - - loss = None - if labels is not None: - shift_logits = logits - if not hasattr(self, "extra_ignored_labels"): - # Fixes https://github.com/unslothai/unsloth/issues/10 - self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda") - pass - - shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) - loss = fast_cross_entropy_loss( - logits = shift_logits, - labels = shift_labels, - ) - pass - # if labels is not None: - # # Shift so that tokens < n predict n - # shift_logits = logits[..., :-1, :].contiguous() - # shift_labels = labels[..., 1:].contiguous() - # # Flatten the tokens - # loss_fct = torch.nn.CrossEntropyLoss() - # shift_logits = shift_logits.view(-1, self.config.vocab_size) - # shift_labels = shift_labels.view(-1) - # # Enable model parallelism - # shift_labels = shift_labels.to(shift_logits.device) - # loss = loss_fct(shift_logits, shift_labels) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) -pass - - class FastGemmaModel(FastLlamaModel): @staticmethod @@ -631,7 +243,7 @@ class FastGemmaModel(FastLlamaModel): # Downcast RoPE embedding to correct data type if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ and (module.cos_cached.dtype != correct_dtype): - + module.cos_cached = module.cos_cached.to(correct_dtype) module.sin_cached = module.sin_cached.to(correct_dtype) pass