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