228 lines
8.9 KiB
Python
228 lines
8.9 KiB
Python
|
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
|