File size: 4,171 Bytes
46144df | 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 | """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."""
# For the default ratio, one group is [KDA, KDA, KDA, MLA]. Placing a
# global-attention layer at the end of an incomplete group avoids ending
# a non-standard-depth model with only recurrent local-state updates.
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":
# Hub config files include descriptive keys such as ``architectures``
# and ``model_type``. Ignoring unknown keys keeps the plain dataclass
# loader compatible with those metadata-rich JSON files.
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)
|