| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class LizzyConfig(PretrainedConfig): |
| model_type = "lizzy" |
| keys_to_ignore_at_inference = ["past_key_values"] |
| base_model_tp_plan = { |
| "layers.*.self_attn.q_proj": "colwise", |
| "layers.*.self_attn.k_proj": "colwise", |
| "layers.*.self_attn.v_proj": "colwise", |
| "layers.*.self_attn.o_proj": "rowwise", |
| "layers.*.mlp.up_proj": "colwise", |
| "layers.*.mlp.gate_proj": "colwise", |
| "layers.*.mlp.down_proj": "rowwise", |
| "lm_head": "colwise", |
| } |
| base_model_pp_plan = { |
| "embed_tokens": (["input_ids"], ["inputs_embeds"]), |
| "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), |
| "norm": (["hidden_states"], ["hidden_states"]), |
| } |
|
|
| def __init__( |
| self, |
| vocab_size: int = 32000, |
| hidden_size: int = 4096, |
| intermediate_size: int = 11008, |
| num_hidden_layers: int = 32, |
| num_attention_heads: int = 32, |
| num_key_value_heads: int | None = None, |
| max_position_embeddings: int = 2048, |
| head_dim: int | None = None, |
| hidden_act: str = "silu", |
| norm_type: str = "rmsnorm", |
| norm_eps: float = 1e-6, |
| norm_has_bias: bool = False, |
| use_pre_attn_norm: bool = True, |
| use_pre_mlp_norm: bool = True, |
| use_post_attn_norm: bool = False, |
| use_post_mlp_norm: bool = False, |
| mlp_type: str = "gated", |
| attention_bias: bool = False, |
| mlp_bias: bool = False, |
| position_embedding_type: str = "rope", |
| rope_theta: float = 10000.0, |
| rope_scaling: dict[str, Any] | None = None, |
| rope_layer_flags: list[bool] | None = None, |
| no_rope_layer_interval: int | None = None, |
| rope_type_overrides: dict[str, str] | None = None, |
| layer_types: list[str] | None = None, |
| layer_layouts: list[str] | None = None, |
| sliding_window: int | None = None, |
| linear_num_key_heads: int | None = None, |
| linear_num_value_heads: int | None = None, |
| linear_key_head_dim: int | None = None, |
| linear_value_head_dim: int | None = None, |
| linear_a_log_min: float | None = None, |
| linear_a_log_max: float | None = None, |
| linear_dt_min: float | None = None, |
| linear_dt_max: float | None = None, |
| linear_dt_init_floor: float | None = None, |
| linear_conv_kernel_dim: int | None = None, |
| linear_allow_neg_eigval: bool | None = None, |
| use_qk_norm: bool = False, |
| qk_norm_type: str = "rmsnorm", |
| attention_dropout: float = 0.0, |
| resid_dropout: float = 0.0, |
| embd_dropout: float = 0.0, |
| initializer_range: float = 0.02, |
| bos_token_id: int | None = None, |
| eos_token_id: int | None = None, |
| pad_token_id: int | None = None, |
| use_cache: bool = True, |
| tie_word_embeddings: bool = False, |
| **kwargs, |
| ) -> None: |
| if num_key_value_heads is None: |
| num_key_value_heads = num_attention_heads |
| if head_dim is None: |
| head_dim = hidden_size // num_attention_heads |
| if no_rope_layer_interval is not None: |
| no_rope_layer_interval = int(no_rope_layer_interval) |
| if no_rope_layer_interval <= 0: |
| no_rope_layer_interval = None |
| if layer_types is None: |
| layer_types = ["full_attention"] * int(num_hidden_layers) |
| if layer_layouts is None: |
| if use_post_attn_norm or use_post_mlp_norm: |
| layer_layouts = ["decoder_postnorm"] * int(num_hidden_layers) |
| else: |
| layer_layouts = ["decoder_prenorm"] * int(num_hidden_layers) |
| if rope_layer_flags is None: |
| rope_enabled = position_embedding_type == "rope" |
| if rope_enabled and no_rope_layer_interval is not None: |
| rope_layer_flags = [ |
| ((layer_idx + 1) % no_rope_layer_interval) != 0 |
| for layer_idx in range(int(num_hidden_layers)) |
| ] |
| else: |
| rope_layer_flags = [rope_enabled] * int(num_hidden_layers) |
|
|
| normalized_rope_scaling = None |
| if rope_scaling is not None: |
| normalized_rope_scaling = dict(rope_scaling) |
| for field_name in ( |
| "factor", |
| "attention_factor", |
| "beta_fast", |
| "beta_slow", |
| ): |
| if normalized_rope_scaling.get(field_name) is not None: |
| normalized_rope_scaling[field_name] = float( |
| normalized_rope_scaling[field_name] |
| ) |
| if ( |
| normalized_rope_scaling.get("original_max_position_embeddings") |
| is not None |
| ): |
| normalized_rope_scaling["original_max_position_embeddings"] = int( |
| normalized_rope_scaling["original_max_position_embeddings"] |
| ) |
|
|
| |
| |
| |
| self.max_position_embeddings = int(max_position_embeddings) |
| self.rope_theta = float(rope_theta) |
| self.rope_scaling = normalized_rope_scaling |
|
|
| super().__init__( |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| pad_token_id=pad_token_id, |
| tie_word_embeddings=tie_word_embeddings, |
| **kwargs, |
| ) |
| self.vocab_size = int(vocab_size) |
| self.hidden_size = int(hidden_size) |
| self.intermediate_size = int(intermediate_size) |
| self.num_hidden_layers = int(num_hidden_layers) |
| self.num_attention_heads = int(num_attention_heads) |
| self.num_key_value_heads = int(num_key_value_heads) |
| self.max_position_embeddings = int(max_position_embeddings) |
| self.head_dim = int(head_dim) |
| self.hidden_act = str(hidden_act) |
| self.norm_type = str(norm_type) |
| self.norm_eps = float(norm_eps) |
| self.norm_has_bias = bool(norm_has_bias) |
| self.use_pre_attn_norm = bool(use_pre_attn_norm) |
| self.use_pre_mlp_norm = bool(use_pre_mlp_norm) |
| self.use_post_attn_norm = bool(use_post_attn_norm) |
| self.use_post_mlp_norm = bool(use_post_mlp_norm) |
| self.mlp_type = str(mlp_type) |
| self.attention_bias = bool(attention_bias) |
| self.mlp_bias = bool(mlp_bias) |
| self.position_embedding_type = str(position_embedding_type) |
| self.rope_theta = float(rope_theta) |
| self.rope_scaling = normalized_rope_scaling |
| self.no_rope_layer_interval = no_rope_layer_interval |
| self.rope_type_overrides = { |
| str(key): str(value) |
| for key, value in dict(rope_type_overrides or {}).items() |
| } |
| self.layer_types = list(layer_types) |
| self.layer_layouts = [str(item) for item in layer_layouts] |
| self.rope_layer_flags = [bool(item) for item in rope_layer_flags] |
| self.sliding_window = sliding_window |
| self.linear_num_key_heads = ( |
| None |
| if linear_num_key_heads is None |
| else int(linear_num_key_heads) |
| ) |
| self.linear_num_value_heads = ( |
| None |
| if linear_num_value_heads is None |
| else int(linear_num_value_heads) |
| ) |
| self.linear_key_head_dim = ( |
| None |
| if linear_key_head_dim is None |
| else int(linear_key_head_dim) |
| ) |
| self.linear_value_head_dim = ( |
| None |
| if linear_value_head_dim is None |
| else int(linear_value_head_dim) |
| ) |
| self.linear_a_log_min = ( |
| None if linear_a_log_min is None else float(linear_a_log_min) |
| ) |
| self.linear_a_log_max = ( |
| None if linear_a_log_max is None else float(linear_a_log_max) |
| ) |
| self.linear_dt_min = ( |
| None if linear_dt_min is None else float(linear_dt_min) |
| ) |
| self.linear_dt_max = ( |
| None if linear_dt_max is None else float(linear_dt_max) |
| ) |
| self.linear_dt_init_floor = ( |
| None |
| if linear_dt_init_floor is None |
| else float(linear_dt_init_floor) |
| ) |
| self.linear_conv_kernel_dim = ( |
| None |
| if linear_conv_kernel_dim is None |
| else int(linear_conv_kernel_dim) |
| ) |
| self.linear_allow_neg_eigval = ( |
| None |
| if linear_allow_neg_eigval is None |
| else bool(linear_allow_neg_eigval) |
| ) |
| self.use_qk_norm = bool(use_qk_norm) |
| self.qk_norm_type = str(qk_norm_type) |
| self.attention_dropout = float(attention_dropout) |
| self.resid_dropout = float(resid_dropout) |
| self.embd_dropout = float(embd_dropout) |
| self.initializer_range = float(initializer_range) |
| self.use_cache = bool(use_cache) |
| self.rms_norm_eps = self.norm_eps |
| self.dtype = None |
|
|
