|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# -*- coding:utf-8 -*- |
| 3 | + |
| 4 | +import math |
| 5 | +from typing import Optional, Tuple |
| 6 | + |
| 7 | +import torch |
| 8 | +from torch import nn |
| 9 | +import transformers |
| 10 | + |
| 11 | + |
| 12 | +def rotate_half(x): |
| 13 | + """Rotates half the hidden dims of the input.""" |
| 14 | + x1 = x[..., : x.shape[-1] // 2].clone() |
| 15 | + x2 = x[..., x.shape[-1] // 2 :].clone() |
| 16 | + return torch.cat((-x2, x1), dim=-1) |
| 17 | + |
| 18 | + |
| 19 | +def apply_rotary_pos_emb(q, k, cos, sin, position_ids): |
| 20 | + gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1] |
| 21 | + gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3]) |
| 22 | + cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) |
| 23 | + sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) |
| 24 | + q_embed = (q * cos) + (rotate_half(q) * sin) |
| 25 | + k_embed = (k * cos) + (rotate_half(k) * sin) |
| 26 | + return q_embed, k_embed |
| 27 | + |
| 28 | + |
| 29 | +def forward( |
| 30 | + self, |
| 31 | + hidden_states: torch.Tensor, |
| 32 | + attention_mask: Optional[torch.Tensor] = None, |
| 33 | + position_ids: Optional[torch.LongTensor] = None, |
| 34 | + past_key_value: Optional[Tuple[torch.Tensor]] = None, |
| 35 | + output_attentions: bool = False, |
| 36 | + use_cache: bool = False, |
| 37 | +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
| 38 | + bsz, q_len, _ = hidden_states.size() |
| 39 | + |
| 40 | + query_states = ( |
| 41 | + self.q_proj(hidden_states) |
| 42 | + .view(bsz, q_len, self.num_heads, self.head_dim) |
| 43 | + .transpose(1, 2) |
| 44 | + ) |
| 45 | + key_states = ( |
| 46 | + self.k_proj(hidden_states) |
| 47 | + .view(bsz, q_len, self.num_heads, self.head_dim) |
| 48 | + .transpose(1, 2) |
| 49 | + ) |
| 50 | + value_states = ( |
| 51 | + self.v_proj(hidden_states) |
| 52 | + .view(bsz, q_len, self.num_heads, self.head_dim) |
| 53 | + .transpose(1, 2) |
| 54 | + ) |
| 55 | + |
| 56 | + kv_seq_len = key_states.shape[-2] |
| 57 | + if past_key_value is not None: |
| 58 | + kv_seq_len += past_key_value[0].shape[-2] |
| 59 | + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) |
| 60 | + query_states, key_states = apply_rotary_pos_emb( |
| 61 | + query_states, key_states, cos, sin, position_ids |
| 62 | + ) |
| 63 | + # [bsz, nh, t, hd] |
| 64 | + |
| 65 | + if past_key_value is not None: |
| 66 | + # reuse k, v, self_attention |
| 67 | + key_states = torch.cat([past_key_value[0], key_states], dim=2) |
| 68 | + value_states = torch.cat([past_key_value[1], value_states], dim=2) |
| 69 | + |
| 70 | + past_key_value = (key_states, value_states) if use_cache else None |
| 71 | + |
| 72 | + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt( |
| 73 | + self.head_dim |
| 74 | + ) |
| 75 | + |
| 76 | + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): |
| 77 | + raise ValueError( |
| 78 | + f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is" |
| 79 | + f" {attn_weights.size()}" |
| 80 | + ) |
| 81 | + |
| 82 | + if attention_mask is not None: |
| 83 | + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): |
| 84 | + raise ValueError( |
| 85 | + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" |
| 86 | + ) |
| 87 | + attn_weights = attn_weights + attention_mask |
| 88 | + attn_weights = torch.max( |
| 89 | + attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min) |
| 90 | + ) |
| 91 | + |
| 92 | + # upcast attention to fp32 |
| 93 | + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( |
| 94 | + query_states.dtype |
| 95 | + ) |
| 96 | + attn_output = torch.matmul(attn_weights, value_states) |
| 97 | + |
| 98 | + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): |
| 99 | + raise ValueError( |
| 100 | + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" |
| 101 | + f" {attn_output.size()}" |
| 102 | + ) |
| 103 | + |
| 104 | + attn_output = attn_output.transpose(1, 2) |
| 105 | + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) |
| 106 | + |
| 107 | + attn_output = self.o_proj(attn_output) |
| 108 | + |
| 109 | + if not output_attentions: |
| 110 | + attn_weights = None |
| 111 | + |
| 112 | + return attn_output, attn_weights, past_key_value |
| 113 | + |
| 114 | + |
| 115 | +def replace_llama_attn_with_non_inplace_operations(): |
| 116 | + """Avoid bugs in mps backend by not using in-place operations.""" |
| 117 | + transformers.models.llama.modeling_llama.LlamaAttention.forward = forward |
| 118 | + |
| 119 | +import transformers |
| 120 | + |
| 121 | + |
| 122 | + |
| 123 | +def replace_llama_attn_with_non_inplace_operations(): |
| 124 | + """Avoid bugs in mps backend by not using in-place operations.""" |
| 125 | + transformers.models.llama.modeling_llama.LlamaAttention.forward = forward |
0 commit comments