configuration_thox_microprocessor.py

"""Hugging Face configuration for THOX Microprocessor (remote code).

This file is uploaded verbatim to the model repository, so it must be
self-contained and must not import the ``thox_microprocessor`` package.
"""

from __future__ import annotations

try:  # transformers >= 5.0 renamed the class; the old name remains an alias.
    from transformers import PreTrainedConfig as _BaseConfig
except ImportError:  # pragma: no cover - transformers 4.x
    from transformers import PretrainedConfig as _BaseConfig


class ThoxMicroprocessorConfig(_BaseConfig):
    model_type = "thox_microprocessor"
    keys_to_ignore_at_inference = ["past_key_values"]
    attribute_map = {
        "hidden_size": "d_model",
        "num_hidden_layers": "num_layers",
        "num_attention_heads": "num_heads",
        "num_key_value_heads": "num_kv_heads",
        "max_position_embeddings": "max_seq_len",
    }

    def __init__(
        self,
        vocab_size: int = 8192,
        d_model: int = 320,
        num_layers: int = 21,
        num_heads: int = 8,
        num_kv_heads: int = 2,
        max_seq_len: int = 1024,
        rope_theta: float = 100000.0,
        rms_eps: float = 1e-6,
        hadamard_block_size: int = 64,
        use_hadamard_mixer: bool = True,
        mixer_center: bool = False,
        mixer_block_shuffle: bool = False,
        use_qk_norm: bool = True,
        use_sandwich_norm: bool = True,
        use_residual_gate: bool = False,
        attention_dropout: float = 0.0,
        embedding_dropout: float = 0.0,
        tie_word_embeddings: bool = True,
        initializer_range: float = 0.02,
        pad_token_id: int = 0,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        architecture_version: str = "thox-microprocessor-v0.2",
        **kwargs,
    ) -> None:
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.max_seq_len = max_seq_len
        self.rope_theta = rope_theta
        self.rms_eps = rms_eps
        self.hadamard_block_size = hadamard_block_size
        self.use_hadamard_mixer = use_hadamard_mixer
        self.mixer_center = mixer_center
        self.mixer_block_shuffle = mixer_block_shuffle
        self.use_qk_norm = use_qk_norm
        self.use_sandwich_norm = use_sandwich_norm
        self.use_residual_gate = use_residual_gate
        self.attention_dropout = attention_dropout
        self.embedding_dropout = embedding_dropout
        self.initializer_range = initializer_range
        self.architecture_version = architecture_version
        # The remote-code reference path recomputes the prefix instead of
        # carrying a KV cache; see modeling_thox_microprocessor for why.
        kwargs.setdefault("use_cache", False)
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
        self.validate()

    @property
    def head_dim(self) -> int:
        return self.d_model // self.num_heads

    @property
    def kv_dim(self) -> int:
        return self.num_kv_heads * self.head_dim

    def validate(self) -> None:
        if self.vocab_size < 256:
            raise ValueError("vocab_size must be >= 256")
        if self.d_model <= 0 or self.num_layers <= 0:
            raise ValueError("d_model and num_layers must be positive")
        if self.num_heads <= 0 or self.num_kv_heads <= 0:
            raise ValueError("num_heads and num_kv_heads must be positive")
        if self.d_model % self.num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")
        if self.num_heads % self.num_kv_heads != 0:
            raise ValueError("num_heads must be divisible by num_kv_heads")
        if self.head_dim % 2 != 0:
            raise ValueError("head_dim must be even for rotary embeddings")
        if self.use_hadamard_mixer:
            if self.hadamard_block_size <= 0 or (
                self.hadamard_block_size & (self.hadamard_block_size - 1)
            ):
                raise ValueError("hadamard_block_size must be a positive power of two")
            if self.d_model % self.hadamard_block_size != 0:
                raise ValueError("d_model must be divisible by hadamard_block_size")
        if not 0.0 <= self.attention_dropout < 1.0:
            raise ValueError("attention_dropout must be in [0, 1)")
        if not 0.0 <= self.embedding_dropout < 1.0:
            raise ValueError("embedding_dropout must be in [0, 1)")


__all__ = ["ThoxMicroprocessorConfig"]