| """Configuration for the OxMini model family.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass, fields |
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import yaml |
|
|
|
|
| @dataclass |
| class OxMiniConfig: |
| """Serializable OxMini architecture configuration. |
| |
| ``use_mhc=False`` selects the hybrid-attention-only ablation. The separate |
| :class:`BaselineGPTForCausalLM` uses the same width/depth fields. |
| """ |
|
|
| vocab_size: int |
| block_size: int = 256 |
| n_layer: int = 8 |
| n_embd: int = 128 |
| n_head: int = 4 |
| kda_to_mla_ratio: int = 3 |
| mla_latent_dim: int = 32 |
| hc_streams: int = 4 |
| use_sinkhorn_mhc: bool = False |
| dropout: float = 0.1 |
| tie_embeddings: bool = True |
| use_mhc: bool = True |
| ffn_multiplier: int = 4 |
| rms_norm_eps: float = 1e-5 |
| bias: bool = False |
| seed: int = 1337 |
|
|
| def __post_init__(self) -> None: |
| if self.vocab_size <= 0: |
| raise ValueError("vocab_size must be positive") |
| if self.block_size <= 0 or self.n_layer <= 0 or self.n_embd <= 0: |
| raise ValueError("block_size, n_layer, and n_embd must be positive") |
| if self.n_head <= 0 or self.n_embd % self.n_head != 0: |
| raise ValueError("n_embd must be divisible by n_head") |
| if (self.n_embd // self.n_head) % 2: |
| raise ValueError("head dimension must be even for RoPE") |
| if not 0 < self.mla_latent_dim < self.n_embd: |
| raise ValueError("mla_latent_dim must be between 0 and n_embd") |
| if self.hc_streams <= 0: |
| raise ValueError("hc_streams must be positive") |
| if self.kda_to_mla_ratio <= 0: |
| raise ValueError("kda_to_mla_ratio must be positive") |
| if not 0.0 <= self.dropout < 1.0: |
| raise ValueError("dropout must be in [0, 1)") |
| if self.ffn_multiplier <= 0: |
| raise ValueError("ffn_multiplier must be positive") |
|
|
| @property |
| def head_dim(self) -> int: |
| return self.n_embd // self.n_head |
|
|
| @property |
| def layer_types(self) -> list[str]: |
| """Return the repeated KDA:MLA pattern, with a final MLA on remainders.""" |
|
|
| |
| |
| |
| group = self.kda_to_mla_ratio + 1 |
| result: list[str] = [] |
| for index in range(self.n_layer): |
| at_group_end = (index + 1) % group == 0 |
| is_final_remainder = index == self.n_layer - 1 and self.n_layer % group != 0 |
| result.append("mla" if at_group_end or is_final_remainder else "kda") |
| return result |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return asdict(self) |
|
|
| @classmethod |
| def from_dict(cls, values: dict[str, Any]) -> "OxMiniConfig": |
| |
| |
| |
| allowed = {field.name for field in fields(cls)} |
| return cls(**{key: value for key, value in values.items() if key in allowed}) |
|
|
| @classmethod |
| def from_file(cls, path: str | Path) -> "OxMiniConfig": |
| path = Path(path) |
| with path.open("r", encoding="utf-8") as handle: |
| if path.suffix.lower() == ".json": |
| values = json.load(handle) |
| else: |
| values = yaml.safe_load(handle) |
| if not isinstance(values, dict): |
| raise ValueError(f"configuration in {path} must be a mapping") |
| return cls.from_dict(values) |
|
|
| def save(self, path: str | Path) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("w", encoding="utf-8") as handle: |
| if path.suffix.lower() == ".json": |
| json.dump(self.to_dict(), handle, indent=2, sort_keys=True) |
| handle.write("\n") |
| else: |
| yaml.safe_dump(self.to_dict(), handle, sort_keys=False) |
|
|
