File size: 3,415 Bytes
8b8e59d | 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 | from dataclasses import dataclass, asdict
@dataclass
class TinyLiquidConfig:
# --- tokenizer ---
vocab_size: int = 8192
# --- architecture (our own "liquid" design, non-transformer) ---
d_model: int = 320 # hidden width
n_blocks: int = 6 # liquid blocks
basis_n: int = 16 # basis blocks in expansion layer
basis_b: int = 4 # dims per basis block (basis_n*basis_b = expanded width)
mlp_ratio: int = 2 # dense MLP hidden = d_model * mlp_ratio
num_experts: int = 0 # 0 = dense gated MLP; >0 = mixture-of-experts MLP
num_experts_per_tok: int = 2
expert_hidden: int = 0 # 0 => mlp_ratio*d_model//2 per expert
num_personas: int = 3 # 0=none, 1=analyst, 2=skeptic (learned persona vectors)
tower_d: int = 0 # 0=off; >0 = wide-head tower width (baseline-preserving growth)
tower_blocks: int = 0 # number of identity-init blocks in the tower
mtp_heads: int = 0 # 0=off; >0 = multi-token prediction aux heads (Meta MTP)
# --- position / normalization ---
rope_theta: float = 10000.0
max_seq_len: int = 1024
norm_eps: float = 1e-6
tie_embeddings: bool = True
def params_estimate(self):
"""Rough parameter count (ignores small heads)."""
d, v = self.d_model, self.vocab_size
n = self.basis_n * self.basis_b
per_block = 2 * n * d # basis: in + forget
if self.num_experts > 0:
h = self.expert_hidden or (self.mlp_ratio * d // 2)
per_mlp = 3 * d * h + d * h + d * self.num_experts
per_block += self.num_experts * per_mlp
else:
h = self.mlp_ratio * d
per_block += 4 * d * h # up, gate, down, forget
params = v * d + self.num_personas * d + self.n_blocks * per_block
if self.tower_d and self.tower_blocks:
td, tn = self.tower_d, self.tower_blocks
tn_ = tn # basis rows in tower
per_tower = 2 * self.basis_n * self.basis_b * td + 4 * (self.mlp_ratio * td) * td
params += td * d + d * td + tn * per_tower
if self.mtp_heads:
params += self.mtp_heads * (d * d + d) # SiLU MLP heads, tied output
return params
CONFIGS = {
"tiny10m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2),
"tiny10m-moe": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
num_experts=4, num_experts_per_tok=2),
"micro6m": dict(d_model=256, n_blocks=5, basis_n=16, basis_b=4, mlp_ratio=2),
"tiny13m": dict(d_model=320, n_blocks=12, basis_n=16, basis_b=4, mlp_ratio=2),
"tiny16m": dict(d_model=448, n_blocks=7, basis_n=16, basis_b=5, mlp_ratio=2),
"tiny20m": dict(d_model=512, n_blocks=8, basis_n=16, basis_b=5, mlp_ratio=2),
"tiny28m": dict(d_model=600, n_blocks=8, basis_n=16, basis_b=6, mlp_ratio=2),
"hybrid18m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
tower_d=512, tower_blocks=4),
"hybrid25m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
tower_d=512, tower_blocks=8),
"hybrid50m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
tower_d=800, tower_blocks=8),
}
|