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,
        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",
        attention_output_gate: bool = False,
        attn_res_block_size: int | None = None,
        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,
        **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.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.attention_output_gate = bool(attention_output_gate)
        self.attn_res_block_size = (
            int(attn_res_block_size) if attn_res_block_size is not None else None
        )
        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.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.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.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")
        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.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)