Blame for tiny_gdn/config.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, fields
from pathlib import Path
from typing import Any

LINEAR_LAYER_TYPES = frozenset({"gdn2", "kda", "diag_kdn", "bigdn2"})
ATTENTION_LAYER_TYPES = frozenset({"full_attention", "gated_mla", "parallax", "bidirectional"})
# Gated-residual token-mixer pairs: Senryu-2's causal one, and Fugue's bidirectional one.
GATED_RESIDUAL_MIXERS = frozenset({("gdn2", "parallax"), ("bigdn2", "bidirectional")})
FUGUE_MODEL_TYPE = "fugue"
RESIDUAL_MODES = frozenset({"pre_norm", "gated_residual"})
DIAG_KDN_HEAD_DIMS = (64, 128)


@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

    # "gated_residual" replaces pre-norm with Qwen's multi-branch gated
    # residual stream: each sublayer reads a gated mix of `residual_branches`
    # normalized branches and writes its output back with per-branch scales.
    residual_mode: str = "pre_norm"
    residual_branches: int = 4
    residual_gate_rank: int = 128
    # GDN-2 output gate as one full-rank projection (Kimi K3) instead of fla's
    # low-rank pair through head_v_dim.
    linear_full_rank_gate: bool = False
    # Hashed n-gram memory (Engram): `ngram_heads_per_order` tables per order,
    # each with a distinct prime row count near `ngram_table_rows`, injected
    # before block `ngram_layer_index`. An empty `ngram_orders` disables it.
    ngram_orders: tuple[int, ...] = ()
    ngram_heads_per_order: int = 4
    ngram_table_rows: int = 2_097_152
    ngram_embedding_dim: int = 128
    ngram_layer_index: int = 1
    ngram_conv_kernel: int = 4

    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
    # Masked-diffusion (Fugue) models replace noised tokens with this id.
    mask_token_id: int | None = None
    # Fugue only: 0 denoises whole sequences bidirectionally (Rondo); G > 0
    # generates G-token blocks left to right over a causal clean stream
    # (Fugue-1, docs/architecture.md "Fugue-1").
    fugue_block_size: int = 0

    def __post_init__(self) -> None:
        for name in ("shared_layer_indices", "ngram_orders"):
            object.__setattr__(self, name, tuple(getattr(self, name)))
        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 LINEAR_LAYER_TYPES:
            raise ValueError(
                f"linear_layer_type must be one of {sorted(LINEAR_LAYER_TYPES)}"
            )
        if self.linear_layer_type == "diag_kdn":
            if self.linear_head_dim not in DIAG_KDN_HEAD_DIMS:
                raise ValueError(
                    f"diag_kdn needs linear_head_dim in {DIAG_KDN_HEAD_DIMS}"
                )
            if self.linear_num_value_heads != self.linear_num_heads:
                raise ValueError("diag_kdn needs linear_num_value_heads == linear_num_heads")
            if self.linear_expand_v != 1.0:
                raise ValueError("diag_kdn needs linear_expand_v == 1.0")
            if self.allow_negative_eigenvalues:
                raise ValueError("diag_kdn derives its gain from the covariance; negative eigenvalues are unsupported")
        if self.attention_layer_type not in ATTENTION_LAYER_TYPES:
            raise ValueError(
                f"attention_layer_type must be one of {sorted(ATTENTION_LAYER_TYPES)}"
            )
        if self.residual_mode not in RESIDUAL_MODES:
            raise ValueError(f"residual_mode must be one of {sorted(RESIDUAL_MODES)}")
        gated = self.residual_mode == "gated_residual"
        if gated:
            if self.residual_branches < 2:
                raise ValueError("gated_residual needs at least two residual_branches")
            if self.residual_gate_rank <= 0:
                raise ValueError("residual_gate_rank must be positive")
            if self.use_block_attn_res:
                raise ValueError("gated_residual replaces block attention residuals")
            if self.mtp_num_heads:
                raise ValueError("gated_residual models do not support MTP heads")
            if self.shared_layer_indices:
                raise ValueError("gated_residual models do not support shared layers")
            if (self.linear_layer_type, self.attention_layer_type) not in GATED_RESIDUAL_MIXERS:
                raise ValueError(
                    "gated_residual models pair gdn2 with parallax attention, or bigdn2 with bidirectional"
                )
        elif (
            self.attention_layer_type in {"parallax", "bidirectional"}
            or self.linear_layer_type == "bigdn2"
            or self.ngram_orders
        ):
            raise ValueError(
                "parallax and bidirectional mixers and n-gram memory need residual_mode 'gated_residual'"
            )
        fugue = self.model_type == FUGUE_MODEL_TYPE
        if fugue != (self.linear_layer_type == "bigdn2"):
            raise ValueError("Fugue models, and only Fugue models, use bidirectional bigdn2 mixers")
        if fugue:
            if self.mask_token_id is None or not 0 <= self.mask_token_id < self.vocab_size:
                raise ValueError("Fugue models need a mask_token_id inside the vocabulary")
            if self.ngram_orders:
                raise ValueError("Fugue models have no n-gram memory")
            if self.fugue_block_size < 0:
                raise ValueError("fugue_block_size cannot be negative")
            if self.fugue_block_size and self.training_sequence_length % self.fugue_block_size:
                raise ValueError("training_sequence_length must be a multiple of fugue_block_size")
        elif self.fugue_block_size:
            raise ValueError("fugue_block_size applies to Fugue models only")
        if self.linear_full_rank_gate and self.linear_layer_type not in {"gdn2", "bigdn2"}:
            raise ValueError("linear_full_rank_gate applies to gdn2 and bigdn2 layers")
        if self.ngram_orders:
            if any(order < 2 for order in self.ngram_orders):
                raise ValueError("ngram_orders must be at least 2")
            if tuple(sorted(set(self.ngram_orders))) != tuple(self.ngram_orders):
                raise ValueError("ngram_orders must be unique and ascending")
            if self.ngram_heads_per_order <= 0 or self.ngram_embedding_dim <= 0:
                raise ValueError("n-gram heads and embedding width must be positive")
            if self.ngram_table_rows <= self.vocab_size:
                raise ValueError("ngram_table_rows must exceed the vocabulary")
            if not 0 <= self.ngram_layer_index < self.num_hidden_layers:
                raise ValueError("ngram_layer_index must name a block")
            if self.ngram_conv_kernel <= 0:
                raise ValueError("ngram_conv_kernel must be positive")
        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)

    @property
    def ngram_table_count(self) -> int:
        return len(self.ngram_orders) * self.ngram_heads_per_order

    @property
    def ngram_memory_dim(self) -> int:
        return self.ngram_table_count * self.ngram_embedding_dim

    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)
        for name in ("shared_layer_indices", "ngram_orders"):
            if name in payload:
                payload[name] = tuple(payload[name])
        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)


def haiku_kdn_config(**overrides: Any) -> TinyGDNConfig:
    """Haiku with Diagonal KDN in place of KDA in the 27 linear layers.

    Opt-in: at equal tokens it trailed KDA in Haiku pilots and trains ~23%
    slower at full depth (docs/architecture.md, "Haiku v2").
    """

    return haiku_config(**{"linear_layer_type": "diag_kdn", **overrides})


def senryu_config(**overrides: Any) -> TinyGDNConfig:
    """800M Senryu: 3×Gated DeltaNet-2 + 1×Gated-MLA over 44 layers, AttnRes, SiTU-GLU.

    Haiku's width, heads, tokenizer, and residual design with GDN-2's
    channel-wise erase and write gates in the 33 linear layers. A rank-512
    query LoRA costs as many parameters as a full 1024×1024 query, so queries
    are full rank, and the KV latent doubles to 512. The feed-forward narrows
    to 3392 to pay for the depth. Next-token only: MTP has no pilot win at
    this scale (docs/architecture.md, "Senryu").
    """

    return haiku_config(
        **{
            "architecture": "SenryuForCausalLM",
            "model_type": "senryu",
            "intermediate_size": 3392,
            "num_hidden_layers": 44,
            "linear_layer_type": "gdn2",
            "mla_q_lora_rank": None,
            "mla_kv_lora_rank": 512,
            "mtp_num_heads": 0,
            "mtp_loss_weight": 0.0,
            **overrides,
        }
    )


def senryu2_1b_config(**overrides: Any) -> TinyGDNConfig:
    """Senryu-2-1B: a 1.09B Senryu-2 backbone plus a 2.2B host-side n-gram memory.

    28 layers of 3×Gated DeltaNet-2 (full-rank output gate) + 1×Parallax
    attention, width 1536, SiTU-GLU 3584, a 4-branch gated residual stream,
    and orders-2/3 hashed n-gram tables injected before block 1
    (docs/architecture.md, "Senryu-2-1B").
    """

    payload: dict[str, Any] = {
        "architecture": "Senryu2ForCausalLM",
        "model_type": "senryu2",
        "vocab_size": 65_536,
        "hidden_size": 1536,
        "intermediate_size": 3584,
        "num_hidden_layers": 28,
        "num_attention_heads": 12,
        "num_key_value_heads": 4,
        "attention_head_dim": 128,
        "full_attention_interval": 4,
        "partial_rotary_factor": 0.5,
        "rope_theta": 1_000_000.0,
        "linear_num_heads": 12,
        "linear_num_value_heads": 12,
        "linear_head_dim": 128,
        "linear_layer_type": "gdn2",
        "attention_layer_type": "parallax",
        "mlp_activation": "situ_glu",
        "residual_mode": "gated_residual",
        "residual_branches": 4,
        "residual_gate_rank": 128,
        "linear_full_rank_gate": True,
        "ngram_orders": (2, 3),
        "ngram_heads_per_order": 4,
        "ngram_table_rows": 2_097_152,
        "ngram_embedding_dim": 128,
        "ngram_layer_index": 1,
        "ngram_conv_kernel": 4,
        "max_position_embeddings": 32_768,
        "training_sequence_length": 32_768,
        "mtp_num_heads": 0,
        "mtp_loss_weight": 0.0,
    }
    payload.update(overrides)
    return TinyGDNConfig(**payload)


def monody_config(**overrides: Any) -> TinyGDNConfig:
    """Monody: the autoregressive control in the Rondo/Monody data-efficiency experiment.

    Senryu-2's backbone at width 1024 with no n-gram memory: 16 layers of
    3×Gated DeltaNet-2 + 1×Parallax attention, SiTU-GLU 2816, 8 query and 4 KV
    heads, trained on 2,048-token windows (docs/architecture.md, "Fugue").
    """

    return senryu2_1b_config(
        **{
            "hidden_size": 1024,
            "intermediate_size": 2816,
            "num_hidden_layers": 16,
            "num_attention_heads": 8,
            "num_key_value_heads": 4,
            "linear_num_heads": 8,
            "linear_num_value_heads": 8,
            "ngram_orders": (),
            "max_position_embeddings": 8192,
            "training_sequence_length": 2048,
            **overrides,
        }
    )


def rondo_config(**overrides: Any) -> TinyGDNConfig:
    """Rondo: the first Fugue model, a masked-diffusion LM on Monody's backbone.

    Monody's widths, depth and heads with every token mixer bidirectional: 12
    bidirectional Gated DeltaNet-2 layers and 4 bidirectional softmax attention
    layers, denoising `<|mask|>` (tokenizer v2, id 22). The backward scans'
    own convolutions and decay projections leave it within 0.3% of Monody's
    parameter count (docs/architecture.md, "Fugue").
    """

    return monody_config(
        **{
            "architecture": "FugueForMaskedDiffusion",
            "model_type": FUGUE_MODEL_TYPE,
            "linear_layer_type": "bigdn2",
            "attention_layer_type": "bidirectional",
            "mask_token_id": 22,
            **overrides,
        }
    )


def fugue1_1b_config(**overrides: Any) -> TinyGDNConfig:
    """Fugue-1-1B: a 1.1B block-diffusion LM on Senryu-2-1B's backbone.

    Senryu-2-1B's widths, depth and heads without n-gram memory, as a Fugue
    model that writes 32-token blocks left to right: a causal clean stream
    (next-token loss, and the exact cache every block reads) and a noisy stream
    denoising each block bidirectionally against that cache, trained on
    4,096-token windows (docs/architecture.md, "Fugue-1").
    """

    return senryu2_1b_config(
        **{
            "architecture": "FugueForMaskedDiffusion",
            "model_type": FUGUE_MODEL_TYPE,
            "linear_layer_type": "bigdn2",
            "attention_layer_type": "bidirectional",
            "mask_token_id": 22,
            "fugue_block_size": 32,
            "ngram_orders": (),
            "max_position_embeddings": 32_768,
            "training_sequence_length": 4096,
            **overrides,
        }
    )


def prompttercet_config(**overrides: Any) -> TinyGDNConfig:
    """PromptTercet: a 501M Senryu-2 prompter, which writes the user's side of a chat.

    Senryu-2-1B's layout at width 1024: 24 layers of 3×Gated DeltaNet-2 +
    1×Parallax attention, SiTU-GLU 2816, 8 query and 4 KV heads, and n-gram
    tables of 1,048,576 rows, 1.09B host-side parameters, half of Senryu-2-1B's
    (docs/architecture.md, "PromptTercet").
    """

    return senryu2_1b_config(
        **{
            "hidden_size": 1024,
            "intermediate_size": 2816,
            "num_hidden_layers": 24,
            "num_attention_heads": 8,
            "num_key_value_heads": 4,
            "linear_num_heads": 8,
            "linear_num_value_heads": 8,
            "ngram_table_rows": 1_048_576,
            **overrides,
        }
    )