swgoo/pmnet

1from typing import Optional2from transformers.modeling_rope_utils import rope_config_validation3from transformers.configuration_utils import layer_type_validation4from transformers.utils import logging5from transformers import PretrainedConfig6 7import transformers.configuration_utils as configuration_util8 9 10logger = logging.get_logger(__name__)11 12 13class PMNetConfig(PretrainedConfig):14 15    model_type = "PMNet"16    keys_to_ignore_at_inference = ["past_key_values"]17 18    base_model_tp_plan = {19        "layers.*.self_attn.q_proj": "colwise",20        "layers.*.self_attn.k_proj": "colwise",21        "layers.*.self_attn.v_proj": "colwise",22        "layers.*.self_attn.o_proj": "rowwise",23        "layers.*.mlp.gate_proj": "colwise",24        "layers.*.mlp.up_proj": "colwise",25        "layers.*.mlp.down_proj": "rowwise",26    }27    base_model_pp_plan = {28        "embed_tokens": (["input_ids"], ["inputs_embeds"]),29        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),30        "norm": (["hidden_states"], ["hidden_states"]),31    }32 33    def __init__(34        self,35        vocab_size: Optional[int] = 151936,36        hidden_size: Optional[int] = 4096,37        intermediate_size: Optional[int] = 22016,38        num_hidden_layers: Optional[int] = 32,39        num_attention_heads: Optional[int] = 32,40        num_key_value_heads: Optional[int] = 32,41        head_dim: Optional[int] = 128,42        memory_size: Optional[int] = 64,43        num_memory: Optional[int] = 32,44        num_memory_read_heads: Optional[int] = 8,45        memory_write_period: Optional[int] = 4,46        hidden_act: Optional[str] = "silu",47        max_position_embeddings: Optional[int] = 32768,48        initializer_range: Optional[float] = 0.02,49        rms_norm_eps: Optional[int] = 1e-6,50        use_cache: Optional[bool] = True,51        tie_word_embeddings: Optional[bool] = False,52        rope_theta=10000.0,53        rope_scaling=None,54        attention_bias: Optional[bool] = False,55        use_sliding_window: Optional[bool] = False,56        sliding_window: Optional[int] = 4096,57        max_window_layers: Optional[int] = 28,58        layer_types: Optional[list[str]] = None,59        attention_dropout: Optional[float] = 0.0,60        memory_cumsum: bool = True,61        **kwargs,62    ):63        self.vocab_size = vocab_size64        self.max_position_embeddings = max_position_embeddings65        self.hidden_size = hidden_size66        self.intermediate_size = intermediate_size67        self.num_hidden_layers = num_hidden_layers68        self.num_attention_heads = num_attention_heads69        self.use_sliding_window = use_sliding_window70        self.sliding_window = sliding_window if self.use_sliding_window else None71        self.max_window_layers = max_window_layers72 73        self.memory_size = memory_size74        self.num_memory = num_memory75        self.num_memory_read_heads = num_memory_read_heads76        self.memory_write_period = memory_write_period77 78        # for backward compatibility79        if num_key_value_heads is None:80            num_key_value_heads = num_attention_heads81 82        self.num_key_value_heads = num_key_value_heads83        self.head_dim = head_dim84        self.hidden_act = hidden_act85        self.initializer_range = initializer_range86        self.rms_norm_eps = rms_norm_eps87        self.use_cache = use_cache88        self.rope_theta = rope_theta89        self.rope_scaling = rope_scaling90        self.attention_bias = attention_bias91        self.attention_dropout = attention_dropout92        self.memory_cumsum = memory_cumsum93 94        if self.rope_scaling is not None and "type" in self.rope_scaling:95            self.rope_scaling["rope_type"] = self.rope_scaling["type"]96        rope_config_validation(self)97 98        self.layer_types = layer_types99        if self.layer_types is None:100            self.layer_types = [101                (102                    "sliding_attention"103                    if self.sliding_window is not None and i >= self.max_window_layers104                    else "full_attention"105                )106                for i in range(self.num_hidden_layers)107            ]108        layer_type_validation(self.layer_types, self.num_hidden_layers)109 110        super().__init__(111            tie_word_embeddings=tie_word_embeddings,112            **kwargs,113        )114 115 116__all__ = ["PMNetConfig"]117