| 708 | |
| 709 | |
| 710 | class LlamaDecoderLayer(nn.Module): |
| 711 | def __init__(self, config: LlamaConfig, layer_idx: int): |
| 712 | super().__init__() |
| 713 | self.hidden_size = config.hidden_size |
| 714 | |
| 715 | config._attn_implementation = "longwriter" |
| 716 | self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) |
| 717 | # print(config._attn_implementation) |
| 718 | |
| 719 | self.mlp = LlamaMLP(config) |
| 720 | self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 721 | self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 722 | |
| 723 | def forward( |
| 724 | self, |
| 725 | hidden_states: torch.Tensor, |
| 726 | attention_mask: Optional[torch.Tensor] = None, |
| 727 | position_ids: Optional[torch.LongTensor] = None, |
| 728 | past_key_value: Optional[Cache] = None, |
| 729 | output_attentions: Optional[bool] = False, |
| 730 | use_cache: Optional[bool] = False, |
| 731 | cache_position: Optional[torch.LongTensor] = None, |
| 732 | position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 |
| 733 | **kwargs, |
| 734 | ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: |
| 735 | """ |
| 736 | Args: |
| 737 | hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` |
| 738 | attention_mask (`torch.FloatTensor`, *optional*): |
| 739 | attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, |
| 740 | query_sequence_length, key_sequence_length)` if default attention is used. |
| 741 | output_attentions (`bool`, *optional*): |
| 742 | Whether or not to return the attentions tensors of all attention layers. See `attentions` under |
| 743 | returned tensors for more detail. |
| 744 | use_cache (`bool`, *optional*): |
| 745 | If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding |
| 746 | (see `past_key_values`). |
| 747 | past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states |
| 748 | cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): |
| 749 | Indices depicting the position of the input sequence tokens in the sequence |
| 750 | position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): |
| 751 | Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, |
| 752 | with `head_dim` being the embedding dimension of each attention head. |
| 753 | kwargs (`dict`, *optional*): |
| 754 | Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code |
| 755 | into the model |
| 756 | """ |
| 757 | residual = hidden_states |
| 758 | |
| 759 | hidden_states = self.input_layernorm(hidden_states) |
| 760 | |
| 761 | # Self Attention |
| 762 | hidden_states, self_attn_weights, present_key_value = self.self_attn( |
| 763 | hidden_states=hidden_states, |
| 764 | attention_mask=attention_mask, |
| 765 | position_ids=position_ids, |
| 766 | past_key_value=past_key_value, |
| 767 | output_attentions=output_attentions, |