jradchenko/DeciCoder-1b

1# coding=utf-82# Copyright and license here3""" PyTorch DeciCoder model."""4import math5from typing import Optional, Tuple6 7import torch8import torch.nn.functional as F9import torch.utils.checkpoint10from torch import nn11from packaging import version12import transformers13if version.parse(transformers.__version__) < version.parse("4.31.0"):14    raise ImportError(15        f"You are using transformers=={transformers.__version__}, but transformers>=4.31.0 is required to use DeciCoder. Please upgrade transformers."16    )17from transformers.models.llama.modeling_llama import LlamaMLP, LlamaRMSNorm, LlamaAttention, apply_rotary_pos_emb, \18    repeat_kv, LlamaPreTrainedModel, LLAMA_START_DOCSTRING, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel19from transformers.utils import add_start_docstrings20 21from .configuration_decicoder import DeciCoderConfig22 23_CONFIG_FOR_DOC = "DeciCoderConfig"24 25 26class DeciCoderAttention(LlamaAttention):27    """Multi-headed attention from 'Attention Is All You Need' paper"""28 29    def __init__(self, config: DeciCoderConfig):30        nn.Module.__init__(self)31        self.config = config32        self.hidden_size = config.hidden_size33        self.num_heads = config.num_attention_heads34        self.head_dim = self.hidden_size // self.num_heads35        self.num_key_value_heads = config.num_key_value_heads36        self.num_key_value_groups = self.num_heads // self.num_key_value_heads37        self.pretraining_tp = config.pretraining_tp38        self.max_position_embeddings = config.max_position_embeddings39        self.rope_theta = getattr(config, 'rope_theta', None)40 41        if (self.head_dim * self.num_heads) != self.hidden_size:42            raise ValueError(43                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"44                f" and `num_heads`: {self.num_heads})."45            )46        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)47        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)48        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)49        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)50 51        self.naive_attention_prefill = config.naive_attention_prefill52        self.naive_attention_decode_batched = config.naive_attention_decode_batched53        self.naive_attention_decode_single = config.naive_attention_decode_single54        self._init_rope()55 56    def forward(57            self,58            hidden_states: torch.Tensor,59            attention_mask: Optional[torch.Tensor] = None,60            position_ids: Optional[torch.LongTensor] = None,61            past_key_value: Optional[Tuple[torch.Tensor]] = None,62            output_attentions: bool = False,63            use_cache: bool = False,64    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:65        bsz, q_len, _ = hidden_states.size()66        if past_key_value is None:67            is_decode = False68        else:69            is_decode = True70        if self.pretraining_tp > 1:71            key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.pretraining_tp72            query_slices = self.q_proj.weight.split((self.num_heads * self.head_dim) // self.pretraining_tp, dim=0)73            key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)74            value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)75 76            query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.pretraining_tp)]77            query_states = torch.cat(query_states, dim=-1)78 79            key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.pretraining_tp)]80            key_states = torch.cat(key_states, dim=-1)81 82            value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.pretraining_tp)]83            value_states = torch.cat(value_states, dim=-1)84 85        else:86            query_states = self.q_proj(hidden_states)87            key_states = self.k_proj(hidden_states)88            value_states = self.v_proj(hidden_states)89 90        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)91        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)92        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)93 94        kv_seq_len = key_states.shape[-2]95        if past_key_value is not None:96            kv_seq_len += past_key_value[0].shape[-2]97        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)98 99        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)100 101        if past_key_value is not None:102            # reuse k, v, self_attention103            key_states = torch.cat([past_key_value[0], key_states], dim=2)104            value_states = torch.cat([past_key_value[1], value_states], dim=2)105 106        past_key_value = (key_states, value_states) if use_cache else None107 108        # repeat k/v heads if n_kv_heads < n_heads109        if is_decode:110            query_states = query_states.view(bsz, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)111            if self.naive_attention_decode_batched and bsz > 1 or self.naive_attention_decode_single and bsz == 1:112                attn_weights = (query_states @ key_states.transpose(-2, -1)) / math.sqrt(key_states.size(-1))113                attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)114                if attention_mask is not None:115                    if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):116                        raise ValueError(117                            f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"118                        )119                    attn_weights = attn_weights + attention_mask120 121                attn_output = torch.matmul(attn_weights, value_states)122            else:123                attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, is_causal=False,124                                                             dropout_p=0.0)125            attn_output = attn_output.contiguous().view(bsz, q_len, self.hidden_size)126 127        else:128            key_states = repeat_kv(key_states, self.num_key_value_groups)129            value_states = repeat_kv(value_states, self.num_key_value_groups)130 131            if not self.naive_attention_prefill:132                attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, is_causal=True,133                                                             dropout_p=0.0)134            else:135                attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)136                # attn_weights = (query_states @ key_states.transpose(-2, -1)) / math.sqrt(key_states.size(-1))137                if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):138                    raise ValueError(139                        f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"140                        f" {attn_weights.size()}"141                    )142 143                if attention_mask is not None:144                    if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):145                        raise ValueError(146                            f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"147                        )148                    attn_weights = attn_weights + attention_mask149 150                # upcast attention to fp32151                attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)152                attn_output = torch.matmul(attn_weights, value_states)153 154            if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):155                raise ValueError(156                    f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"157                    f" {attn_output.size()}"158                )159 160            attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)161            # attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)162 163        if self.pretraining_tp > 1:164            attn_output = attn_output.split(self.hidden_size // self.pretraining_tp, dim=2)165            o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.pretraining_tp, dim=1)166            attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.pretraining_tp)])167        else:168            attn_output = self.o_proj(attn_output)169 170        if not output_attentions:171            attn_weights = None172 173        return attn_output, attn_weights, past_key_value174 175 176class DeciCoderDecoderLayer(LlamaDecoderLayer):177    def __init__(self, config: DeciCoderConfig):178        nn.Module.__init__(self)179        self.hidden_size = config.hidden_size180        self.self_attn = DeciCoderAttention(config=config)181        self.mlp = LlamaMLP(config)182        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)183        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)184 185 186@add_start_docstrings(187    "The bare DeciCoder Model outputting raw hidden-states without any specific head on top.",188    LLAMA_START_DOCSTRING,189)190class DeciCoderPreTrainedModel(LlamaPreTrainedModel):191    config_class = DeciCoderConfig192    _no_split_modules = ["DeciCoderDecoderLayer"]193    _keys_to_ignore_on_load_missing = ["self_attn.rotary_emb.inv_freq"]194 195 196@add_start_docstrings(197    "The bare DeciCoder Model outputting raw hidden-states without any specific head on top.",198    LLAMA_START_DOCSTRING,199)200class DeciCoderModel(LlamaModel, DeciCoderPreTrainedModel):201    """202    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeciCoderDecoderLayer`]203 204    Args:205        config: DeciCoderConfig206    """207 208    def __init__(self, config: DeciCoderConfig):209        DeciCoderPreTrainedModel.__init__(self, config)210        self.padding_idx = config.pad_token_id211        self.vocab_size = config.vocab_size212 213        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)214        self.layers = nn.ModuleList([DeciCoderDecoderLayer(config) for _ in range(config.num_hidden_layers)])215        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)216 217        self.gradient_checkpointing = False218        # Initialize weights and apply final processing219        self.post_init()220 221    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):222        self._validate_config_supports_attention_mask(attention_mask, input_shape, past_key_values_length)223        return LlamaModel._prepare_decoder_attention_mask(224            self, attention_mask, input_shape, inputs_embeds, past_key_values_length)225 226    def _validate_config_supports_attention_mask(self, attention_mask, input_shape, past_key_values_length):227        is_decode = past_key_values_length > 0228        if not torch.all(torch.eq(attention_mask, 1)).item():229            if is_decode:230                if input_shape[0] == 1 and not self.config.naive_attention_decode_single:231                    raise ValueError(232                        "For support of custom attention masks please set naive_attention_decode_single to True in the "233                        "config")234                elif input_shape[0] > 1 and not self.config.naive_attention_decode_batched:235                    raise ValueError(236                        "For support of custom attention masks please set naive_attention_decode_batched to True in the"237                        "config")238            else:239                if not self.config.naive_attention_prefill:240                    raise ValueError("For support of custom attention masks please set naive_attention_prefill to "241                                     "True in the config")242 243 244class DeciCoderForCausalLM(LlamaForCausalLM, DeciCoderPreTrainedModel):245    def __init__(self, config):246        DeciCoderPreTrainedModel.__init__(self, config)247        self.model = DeciCoderModel(config)248        self.pretraining_tp = config.pretraining_tp249        self.vocab_size = config.vocab_size250        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)251 252        # Initialize weights and apply final processing253        self.post_init()254