File size: 8,042 Bytes
d03200e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """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)
|