File size: 3,270 Bytes
dc22a00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
from transformers.models.llama.modeling_llama import *
from beagle.mixin import *
class LlamaBeagleAttention(LlamaAttention, BeagleAttentionMixin):
def forward(
self,
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,
use_cache: bool = False,
cache_position: Optional[torch.LongTensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
bsz, q_len, _ = hidden_states.size()
query_states, key_states, value_states = self.qkv_transform(
hidden_states, past_key_value, use_cache, position_embeddings, **kwargs)
################################################
### everything kept original starting from here
################################################
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
if attention_mask is not None: # no matter the length, we just slice it
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
attn_weights = attn_weights + causal_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)
if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
raise ValueError(
f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
f" {attn_output.size()}"
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, -1)
attn_output = self.o_proj(attn_output)
if not output_attentions:
attn_weights = None
return attn_output, attn_weights, past_key_value
class LlamaBeagleDecoderLayer(LlamaDecoderLayer):
def __init__(self, config, layer_id):
super().__init__(config, layer_id)
if not config.beagle_use_fc_eagle:
delattr(self, 'self_attn')
recycle_vram()
self.self_attn = LlamaBeagleAttention(
config=config, layer_idx=0
)
class LlamaForSpeculativeCausalLM(LlamaForCausalLM, BeagleMixin):
_no_split_modules = ["LlamaDecoderLayer", "LlamaBeagleDecoderLayer"]
def __init__(self, config):
super().__init__(config)
BeagleMixin.__init__(self, config)
self.speculative_decoder = LlamaBeagleDecoderLayer(config, layer_id=0)
self.post_init()
def forward(self, *args, **kwargs) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
return self.beagle_forward(*args, **kwargs)
|