configuration_microloop.py

"""Configuration for the MicroLoop-Diffusion model.

The configuration is intentionally explicit.  It is the single source of truth for
the parameter-count gate and is serializable by Hugging Face when Transformers is
installed.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import yaml

try:  # Keep config inspection useful before optional HF integration is installed.
    from transformers import PretrainedConfig
except ImportError:  # pragma: no cover - exercised only in a minimal environment.

    class PretrainedConfig:  # type: ignore[no-redef]
        model_type = "microloop_diffusion"

        def __init__(self, **kwargs: Any) -> None:
            for key, value in kwargs.items():
                setattr(self, key, value)

        def to_dict(self) -> dict[str, Any]:
            return dict(self.__dict__)


class MicroLoopConfig(PretrainedConfig):
    """Model, diffusion, and selective-looping configuration.

    The defaults match the locked 10M specification.  Feature configuration is
    stored on the model config for deterministic HF save/reload and is also emitted
    separately as ``diffusion_config.json`` by the eventual release exporter.
    """

    model_type = "microloop_diffusion"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size: int = 8192,
        hidden_size: int = 240,
        num_hidden_layers: int = 12,
        num_attention_heads: int = 6,
        num_key_value_heads: int = 2,
        head_dimension: int = 40,
        intermediate_size: int = 640,
        ffn_rank: int | None = None,
        ffn_factor_activation: str = "silu",
        activation: str = "swiglu",
        normalization: str = "rmsnorm",
        positional_encoding: str = "rope",
        tie_word_embeddings: bool = True,
        max_position_embeddings: int = 2048,
        dropout: float = 0.0,
        attention_implementation: str = "eager",
        qk_norm: str = "none",
        qk_norm_position: str = "pre_rope",
        attention_output_gate: bool = False,
        attention_output_gate_activation: str = "silu",
        attn_res_block_size: int | None = None,
        mhc_multiplier: int = 1,
        mhc_sinkhorn_iterations: int = 20,
        mhc_eps: float = 1e-6,
        mhc_init_scale: float = 0.01,
        mtp_enabled: bool = False,
        swiglu_clamp: dict[str, Any] | None = None,
        rms_norm_eps: float = 1e-5,
        rope_theta: float = 10000.0,
        architecture: str = "MicroLoopForDiffusionLM",
        target_parameters: int = 10_000_000,
        diffusion: dict[str, Any] | None = None,
        looping: dict[str, Any] | None = None,
        tokenizer: dict[str, Any] | None = None,
        digit_position_embedding: dict[str, Any] | None = None,
        ngram_memory: dict[str, Any] | None = None,
        value_residual: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        kwargs.setdefault("is_decoder", True)
        kwargs.setdefault("is_encoder_decoder", False)
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
        self.vocab_size = int(vocab_size)
        self.hidden_size = int(hidden_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.head_dimension = int(head_dimension)
        self.intermediate_size = int(intermediate_size)
        self.ffn_rank = ffn_rank
        self.ffn_factor_activation = str(ffn_factor_activation)
        self.activation = activation
        self.normalization = normalization
        self.positional_encoding = positional_encoding
        self.tie_word_embeddings = bool(tie_word_embeddings)
        self.max_position_embeddings = int(max_position_embeddings)
        self.dropout = float(dropout)
        self.attention_implementation = str(attention_implementation)
        self.qk_norm = str(qk_norm)
        self.qk_norm_position = str(qk_norm_position)
        self.attention_output_gate = bool(attention_output_gate)
        self.attention_output_gate_activation = str(attention_output_gate_activation)
        self.attn_res_block_size = (
            int(attn_res_block_size) if attn_res_block_size is not None else None
        )
        self.mhc_multiplier = int(mhc_multiplier)
        self.mhc_sinkhorn_iterations = int(mhc_sinkhorn_iterations)
        self.mhc_eps = float(mhc_eps)
        self.mhc_init_scale = float(mhc_init_scale)
        self.mtp_enabled = bool(mtp_enabled)
        self.swiglu_clamp = dict(swiglu_clamp or {})
        self.rms_norm_eps = float(rms_norm_eps)
        self.rope_theta = float(rope_theta)
        self.architecture = architecture
        self.target_parameters = int(target_parameters)
        self.diffusion = dict(diffusion or {})
        self.looping = dict(looping or {})
        self.tokenizer = dict(tokenizer or {})
        self.digit_position_embedding = dict(digit_position_embedding or {})
        self.ngram_memory = dict(ngram_memory or {})
        self.value_residual = dict(value_residual or {})
        if self.ngram_memory:
            # Unversioned checkpoints were trained with the original linear hash.
            self.ngram_memory.setdefault("hash_version", "legacy_v1")
        self.validate()

    @property
    def head_dim(self) -> int:
        return self.head_dimension

    @classmethod
    def from_yaml(cls, path: str | Path) -> "MicroLoopConfig":
        """Load the locked nested YAML layout used by the project configs."""

        payload = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
        model = dict(payload.get("model", payload))
        model.pop("architecture", None) if model.get("architecture") is None else None
        return cls(
            **model,
            diffusion=payload.get("diffusion", {}),
            looping=payload.get("looping", {}),
            tokenizer=payload.get("tokenizer", {}),
        )

    def validate(self) -> None:
        """Raise a clear error for shape or locked-spec inconsistencies."""

        positive = {
            "vocab_size": self.vocab_size,
            "hidden_size": self.hidden_size,
            "num_hidden_layers": self.num_hidden_layers,
            "num_attention_heads": self.num_attention_heads,
            "num_key_value_heads": self.num_key_value_heads,
            "head_dimension": self.head_dimension,
            "intermediate_size": self.intermediate_size,
            "max_position_embeddings": self.max_position_embeddings,
        }
        invalid = [name for name, value in positive.items() if value <= 0]
        if invalid:
            raise ValueError(f"Configuration values must be positive: {', '.join(invalid)}")
        if self.ffn_rank is not None and (
            type(self.ffn_rank) is not int
            or not 0 < self.ffn_rank <= min(self.hidden_size, self.intermediate_size)
        ):
            raise ValueError(
                "ffn_rank must be an integer in [1, min(hidden_size, intermediate_size)]"
            )
        if self.ffn_factor_activation not in {"silu", "identity"}:
            raise ValueError("ffn_factor_activation must be silu or identity")
        if self.hidden_size != self.num_attention_heads * self.head_dimension:
            raise ValueError(
                "hidden_size must equal num_attention_heads * head_dimension: "
                f"{self.hidden_size} != {self.num_attention_heads} * {self.head_dimension}"
            )
        if self.num_attention_heads % self.num_key_value_heads:
            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
        if self.head_dimension % 2:
            raise ValueError("RoPE requires an even head_dimension")
        if self.dropout < 0.0 or self.dropout >= 1.0:
            raise ValueError("dropout must be in [0, 1)")
        if self.attention_implementation not in {"eager", "sdpa"}:
            raise ValueError("attention_implementation must be eager or sdpa")
        if self.qk_norm not in {"none", "per_head"}:
            raise ValueError("qk_norm must be none or per_head")
        if self.qk_norm_position not in {"pre_rope", "post_rope"}:
            raise ValueError("qk_norm_position must be pre_rope or post_rope")
        if self.attention_output_gate_activation not in {"silu", "sigmoid"}:
            raise ValueError("attention_output_gate_activation must be silu or sigmoid")
        if self.attn_res_block_size is not None and self.attn_res_block_size < 2:
            raise ValueError("attn_res_block_size must be at least two when enabled")
        loop_mode = str(self.looping.get("mode", "layer"))
        if loop_mode not in {"layer", "block"}:
            raise ValueError("looping mode must be layer or block")
        loop_gated = bool(self.looping.get("gated", False))
        if loop_gated and loop_mode != "block":
            raise ValueError("gated looping requires looping mode=block")
        max_loop_count = int(
            self.looping.get("max_loop_count", self.looping.get("maximum_serving_loops", 3))
        )
        if max_loop_count < 1:
            raise ValueError("looping max_loop_count must be positive")
        loop_layers = [int(layer) for layer in self.looping.get("layers", [4, 5, 6])]
        if loop_mode == "block" and loop_layers:
            valid_layers = sorted(
                {layer for layer in loop_layers if 1 <= layer <= self.num_hidden_layers}
            )
            if valid_layers and valid_layers != list(range(valid_layers[0], valid_layers[-1] + 1)):
                raise ValueError("block looping layers must form a contiguous range")
        if self.mhc_multiplier < 1:
            raise ValueError("mhc_multiplier must be at least one")
        if self.mhc_sinkhorn_iterations < 1:
            raise ValueError("mhc_sinkhorn_iterations must be at least one")
        if self.mhc_eps <= 0:
            raise ValueError("mhc_eps must be positive")
        if self.mhc_init_scale <= 0:
            raise ValueError("mhc_init_scale must be positive")
        if self.mhc_multiplier > 1 and self.attn_res_block_size is not None:
            raise ValueError("mHC and attn_res_block_size cannot be enabled together")
        if self.swiglu_clamp:
            enabled = bool(self.swiglu_clamp.get("enabled", False))
            if enabled:
                linear_min = float(self.swiglu_clamp.get("linear_min", -10.0))
                linear_max = float(self.swiglu_clamp.get("linear_max", 10.0))
                gate_max = float(self.swiglu_clamp.get("gate_max", 10.0))
                if linear_min >= linear_max:
                    raise ValueError("swiglu_clamp linear_min must be below linear_max")
                if gate_max <= 0:
                    raise ValueError("swiglu_clamp gate_max must be positive")
        if self.ngram_memory.get("enabled", False):
            if self.ngram_memory["hash_version"] not in {"legacy_v1", "polynomial_v2"}:
                raise ValueError("ngram_memory hash_version must be legacy_v1 or polynomial_v2")
            orders = self.ngram_memory.get("orders", [2, 3])
            if not isinstance(orders, (list, tuple)) or not orders:
                raise ValueError("ngram_memory orders must be a non-empty list")
            if any(not isinstance(order, int) or order < 2 for order in orders):
                raise ValueError("ngram_memory orders must contain integers >= 2")
            if len(set(orders)) != len(orders):
                raise ValueError("ngram_memory orders must be unique")
            if self.ngram_memory.get("mode", "lookup") not in {"lookup", "parameter_free"}:
                raise ValueError("ngram_memory mode must be lookup or parameter_free")
            for name in ("num_hash_heads", "num_buckets", "embedding_dim", "insertion_layer"):
                value = int(self.ngram_memory.get(name, 0))
                if value <= 0:
                    raise ValueError(f"ngram_memory {name} must be positive")
            insertion_layer = int(self.ngram_memory["insertion_layer"])
            if insertion_layer > self.num_hidden_layers:
                raise ValueError("ngram_memory insertion_layer exceeds num_hidden_layers")
            if self.ngram_memory.get("canonicalization", "raw_math_safe") != "raw_math_safe":
                raise ValueError("ngram_memory canonicalization must be raw_math_safe")
        if self.value_residual.get("enabled", False) and self.mhc_multiplier > 1:
            raise ValueError("value_residual is not supported with mHC")
        digit_settings = self.digit_position_embedding
        if digit_settings.get("enabled", False):
            max_positions = int(digit_settings.get("max_positions", 128))
            digit_token_ids = digit_settings.get("digit_token_ids", [])
            if max_positions < 1:
                raise ValueError("digit_position_embedding max_positions must be positive")
            if len(digit_token_ids) != 10 or len(set(digit_token_ids)) != 10:
                raise ValueError(
                    "digit_position_embedding digit_token_ids must contain ten unique IDs"
                )
            if any(
                int(token_id) < 0 or int(token_id) >= self.vocab_size
                for token_id in digit_token_ids
            ):
                raise ValueError("digit_position_embedding digit_token_ids must be in vocabulary")
        if self.activation.lower() != "swiglu":
            raise ValueError("M0 only implements the locked SwiGLU activation")
        if self.normalization.lower() != "rmsnorm":
            raise ValueError("M0 only implements the locked RMSNorm normalization")
        if self.positional_encoding.lower() != "rope":
            raise ValueError("M0 only implements the locked RoPE positional encoding")

    def diffusion_dict(self) -> dict[str, Any]:
        """Return a copy suitable for a standalone diffusion config artifact."""

        return dict(self.diffusion)

    def looping_dict(self) -> dict[str, Any]:
        """Return a copy suitable for experiment logging."""

        return dict(self.looping)