tiny_gdn/config.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
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

    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")

    @property
    def layer_types(self) -> tuple[str, ...]:
        return tuple(
            "full_attention" if (index + 1) % self.full_attention_interval == 0 else "gdn2"
            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"])
        return cls(**payload)