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"] ) # Transformers validates RoPE settings during PretrainedConfig # initialization, so publish the rope-critical fields before # calling `super().__init__()`. 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