from __future__ import annotations import json from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import Any @dataclass(frozen=True) class TinyGDNConfig: architecture: str = "TinyGDNForCausalLM" model_type: str = "tiny_gdn" vocab_size: int = 49_152 # Deep-thin sizing is deliberate: controlled sub-billion studies find # depth materially more valuable than width around the 125M-150M scale. hidden_size: int = 512 intermediate_size: int = 1_472 num_hidden_layers: int = 32 num_attention_heads: int = 4 num_key_value_heads: int = 1 attention_head_dim: int = 128 full_attention_interval: int = 4 attention_dropout: float = 0.0 partial_rotary_factor: float = 0.5 rope_theta: float = 1_000_000.0 linear_num_heads: int = 4 linear_num_value_heads: int = 4 linear_head_dim: int = 128 linear_expand_v: float = 1.0 linear_conv_kernel_dim: int = 4 allow_negative_eigenvalues: bool = False linear_layer_type: str = "gdn2" attention_layer_type: str = "full_attention" mlp_activation: str = "swiglu" use_block_attn_res: bool = False kda_safe_gate: bool = True kda_lower_bound: float = -5.0 mla_q_lora_rank: int | None = 512 mla_kv_lora_rank: int = 256 mla_qk_nope_head_dim: int = 128 mla_v_head_dim: int = 128 situ_gate_cap: float = 4.0 situ_up_cap: float = 25.0 max_position_embeddings: int = 32_768 training_sequence_length: int = 2_048 rms_norm_eps: float = 1e-6 initializer_range: float = 0.02 tie_word_embeddings: bool = True shared_layer_indices: tuple[int, ...] = () # MTP is an opt-in ablation at this scale; static MTP is not assumed to # improve a 150M model without a controlled pilot. mtp_num_heads: int = 0 mtp_adapter_rank: int = 128 mtp_loss_weight: float = 0.0 bos_token_id: int = 0 eos_token_id: int = 1 pad_token_id: int = 2 unk_token_id: int = 3 def __post_init__(self) -> None: if self.vocab_size <= 0 or self.vocab_size > 65_536: raise ValueError("vocab_size must fit the uint16 token dataset") if self.hidden_size != self.num_attention_heads * self.attention_head_dim: raise ValueError("hidden_size must equal num_attention_heads * attention_head_dim") if self.hidden_size != self.linear_num_heads * self.linear_head_dim: raise ValueError("hidden_size must equal linear_num_heads * linear_head_dim") if self.linear_num_value_heads < self.linear_num_heads: raise ValueError("linear_num_value_heads must be at least linear_num_heads") if self.linear_num_value_heads % self.linear_num_heads != 0: raise ValueError("linear_num_value_heads must be divisible by linear_num_heads") if self.num_attention_heads % self.num_key_value_heads != 0: raise ValueError("num_attention_heads must be divisible by num_key_value_heads") if self.num_hidden_layers % self.full_attention_interval != 0: raise ValueError("num_hidden_layers must be divisible by full_attention_interval") if not 0.0 < self.partial_rotary_factor <= 1.0: raise ValueError("partial_rotary_factor must be in (0, 1]") rotary_dim = int(self.attention_head_dim * self.partial_rotary_factor) if rotary_dim <= 0 or rotary_dim % 2: raise ValueError("The partial rotary dimension must be positive and even") if self.training_sequence_length > self.max_position_embeddings: raise ValueError("training_sequence_length exceeds max_position_embeddings") if len(set(self.shared_layer_indices)) != len(self.shared_layer_indices): raise ValueError("shared_layer_indices must be unique") if any( index < 0 or index >= self.num_hidden_layers for index in self.shared_layer_indices ): raise ValueError("shared_layer_indices contains an invalid layer") if self.mtp_num_heads < 0: raise ValueError("mtp_num_heads cannot be negative") if self.mtp_num_heads and self.mtp_adapter_rank <= 0: raise ValueError("mtp_adapter_rank must be positive when MTP is enabled") if not 0.0 <= self.mtp_loss_weight <= 1.0: raise ValueError("mtp_loss_weight must be between zero and one") for token_id in ( self.bos_token_id, self.eos_token_id, self.pad_token_id, self.unk_token_id, ): if not 0 <= token_id < self.vocab_size: raise ValueError(f"Special token ID {token_id} is outside the vocabulary") if self.linear_layer_type not in {"gdn2", "kda"}: raise ValueError("linear_layer_type must be 'gdn2' or 'kda'") if self.attention_layer_type not in {"full_attention", "gated_mla"}: raise ValueError("attention_layer_type must be 'full_attention' or 'gated_mla'") if self.mlp_activation not in {"swiglu", "situ_glu"}: raise ValueError("mlp_activation must be 'swiglu' or 'situ_glu'") if self.mla_kv_lora_rank <= 0: raise ValueError("mla_kv_lora_rank must be positive") if self.mla_q_lora_rank is not None and self.mla_q_lora_rank <= 0: raise ValueError("mla_q_lora_rank must be positive or None") if self.mla_qk_nope_head_dim <= 0 or self.mla_v_head_dim <= 0: raise ValueError("MLA head dimensions must be positive") if self.kda_lower_bound >= 0.0: raise ValueError("kda_lower_bound must be negative log-decay") if self.situ_gate_cap <= 0.0 or self.situ_up_cap <= 0.0: raise ValueError("SiTU caps must be positive") @property def layer_types(self) -> tuple[str, ...]: return tuple( self.attention_layer_type if (index + 1) % self.full_attention_interval == 0 else self.linear_layer_type for index in range(self.num_hidden_layers) ) @property def rotary_dim(self) -> int: return int(self.attention_head_dim * self.partial_rotary_factor) @property def effective_num_layers(self) -> int: return self.num_hidden_layers + len(self.shared_layer_indices) def to_dict(self) -> dict[str, Any]: payload = asdict(self) payload["layer_types"] = list(self.layer_types) return payload def save_json(self, path: Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text( json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8", ) @classmethod def from_json(cls, path: Path) -> TinyGDNConfig: payload = json.loads(path.read_text(encoding="utf-8")) payload.pop("layer_types", None) if "shared_layer_indices" in payload: payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"]) allowed = {item.name for item in fields(cls)} return cls(**{key: value for key, value in payload.items() if key in allowed}) def haiku_config(**overrides: Any) -> TinyGDNConfig: """650M Haiku: 3×KDA + 1×Gated-MLA, AttnRes, SiTU-GLU, MTP.""" payload = { "architecture": "HaikuForCausalLM", "model_type": "haiku", "vocab_size": 65_536, "hidden_size": 1024, "intermediate_size": 3840, "num_hidden_layers": 36, "num_attention_heads": 8, "num_key_value_heads": 8, "attention_head_dim": 128, "full_attention_interval": 4, "linear_num_heads": 8, "linear_num_value_heads": 8, "linear_head_dim": 128, "linear_layer_type": "kda", "attention_layer_type": "gated_mla", "mlp_activation": "situ_glu", "use_block_attn_res": True, "mtp_num_heads": 2, "mtp_adapter_rank": 128, "mtp_loss_weight": 0.2, "max_position_embeddings": 32_768, "training_sequence_length": 2048, } payload.update(overrides) return TinyGDNConfig(**payload)