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