Spaces:
Runtime error
Runtime error
| import math | |
| from typing import Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| from transformers.models.llama.modeling_llama import ( | |
| Cache, | |
| LlamaAttention, | |
| LlamaFlashAttention2, | |
| apply_rotary_pos_emb, | |
| repeat_kv, | |
| ) | |
| from transformers.utils import logging | |
| logger = logging.get_logger(__name__) | |
| # Modified from: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py | |
| def llama_torch_attn_forward( | |
| self: "LlamaAttention", | |
| hidden_states: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| position_ids: Optional[torch.LongTensor] = None, | |
| past_key_value: Optional["Cache"] = None, | |
| output_attentions: bool = False, | |
| **kwargs, | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: | |
| bsz, q_len, _ = hidden_states.size() | |
| query_states = self.q_proj(hidden_states) | |
| key_states = self.k_proj(hidden_states) | |
| value_states = self.v_proj(hidden_states) | |
| query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) | |
| key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
| value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
| kv_seq_len = key_states.shape[-2] | |
| if past_key_value is not None: | |
| kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) | |
| cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) | |
| query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) | |
| if past_key_value is not None: | |
| cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models | |
| key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) | |
| key_states = repeat_kv(key_states, self.num_key_value_groups) | |
| value_states = repeat_kv(value_states, self.num_key_value_groups) | |
| if getattr(self.config, "group_size_ratio", None) and self.training: # shift | |
| groupsz = int(q_len * getattr(self.config, "group_size_ratio")) | |
| assert q_len % groupsz == 0, "q_len {} should be divisible by group size {}.".format(q_len, groupsz) | |
| num_groups = q_len // groupsz | |
| def shift(state: torch.Tensor) -> torch.Tensor: | |
| state = state.transpose(1, 2) # output: (bsz, seq_len, n_heads, head_dim) | |
| state = torch.cat( | |
| (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)), | |
| dim=2, | |
| ) | |
| return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim).transpose(1, 2) | |
| query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states) | |
| if attention_mask is not None: | |
| attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1) | |
| attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) | |
| if attention_mask is not None: | |
| attn_weights = attn_weights + attention_mask | |
| # upcast attention to fp32 | |
| attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) | |
| attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) | |
| attn_output = torch.matmul(attn_weights, value_states) # (bsz, :, seq_len, :) or (bsz*n_group, :, groupsz, :) | |
| attn_output = attn_output.transpose(1, 2).contiguous() | |
| if getattr(self.config, "group_size_ratio", None) and self.training: # shift back | |
| attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim) | |
| attn_output = torch.cat( | |
| ( | |
| attn_output[:, :, : self.num_heads // 2], | |
| attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1), | |
| ) | |
| ) | |
| attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) | |
| attn_output = self.o_proj(attn_output) | |
| if not output_attentions: | |
| attn_weights = None | |
| return attn_output, attn_weights, past_key_value | |
| # Modified from: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py | |
| def llama_flash_attn_forward( | |
| self: "LlamaFlashAttention2", | |
| hidden_states: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| position_ids: Optional[torch.LongTensor] = None, | |
| past_key_value: Optional[Tuple[torch.Tensor]] = None, | |
| output_attentions: bool = False, | |
| **kwargs, | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: | |
| # LlamaFlashAttention2 attention does not support output_attentions | |
| output_attentions = False | |
| bsz, q_len, _ = hidden_states.size() | |
| query_states = self.q_proj(hidden_states) | |
| key_states = self.k_proj(hidden_states) | |
| value_states = self.v_proj(hidden_states) | |
| # FlashAttention requires the input to have the shape (bsz, seq_len, n_heads, head_dim) | |
| query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) | |
| key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
| value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
| kv_seq_len = key_states.shape[-2] | |
| if past_key_value is not None: | |
| kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) | |
| cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) | |
| query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) | |
| if past_key_value is not None: | |
| cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models | |
| key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) | |
| key_states = repeat_kv(key_states, self.num_key_value_groups) | |
| value_states = repeat_kv(value_states, self.num_key_value_groups) | |
| query_states = query_states.transpose(1, 2) # (bsz, seq_len, n_heads, head_dim) | |
| key_states = key_states.transpose(1, 2) # (bsz, seq_len, n_heads, head_dim) | |
| value_states = value_states.transpose(1, 2) # (bsz, seq_len, n_heads, head_dim) | |
| dropout_rate = self.attention_dropout if self.training else 0.0 | |
| input_dtype = query_states.dtype | |
| if input_dtype == torch.float32: | |
| if torch.is_autocast_enabled(): | |
| target_dtype = torch.get_autocast_gpu_dtype() | |
| elif hasattr(self.config, "_pre_quantization_dtype"): | |
| target_dtype = self.config._pre_quantization_dtype | |
| else: | |
| target_dtype = self.q_proj.weight.dtype | |
| logger.warning_once("The input hidden states seems to be silently casted in float32.") | |
| query_states = query_states.to(target_dtype) | |
| key_states = key_states.to(target_dtype) | |
| value_states = value_states.to(target_dtype) | |
| if getattr(self.config, "group_size_ratio", None) and self.training: # shift | |
| groupsz = int(q_len * getattr(self.config, "group_size_ratio")) | |
| assert q_len % groupsz == 0, "q_len {} should be divisible by group size {}.".format(q_len, groupsz) | |
| num_groups = q_len // groupsz | |
| def shift(state: torch.Tensor) -> torch.Tensor: | |
| state = torch.cat( | |
| (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)), | |
| dim=2, | |
| ) | |
| return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim) | |
| query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states) | |
| if attention_mask is not None: | |
| attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1) | |
| attn_output: torch.Tensor = self._flash_attention_forward( | |
| query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate | |
| ) | |
| if getattr(self.config, "group_size_ratio", None) and self.training: # shift back | |
| attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim) | |
| attn_output = torch.cat( | |
| ( | |
| attn_output[:, :, : self.num_heads // 2], | |
| attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1), | |
| ) | |
| ) | |
| attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() | |
| attn_output = self.o_proj(attn_output) | |
| if not output_attentions: | |
| attn_weights = None | |
| return attn_output, attn_weights, past_key_value | |
| def apply_llama_patch() -> None: | |
| LlamaAttention.forward = llama_torch_attn_forward | |
| LlamaFlashAttention2.forward = llama_flash_attn_forward | |