""" Quazimoto-LM — a language model whose channel-mixing block is a bank of coupled phase oscillators (Kuramoto dynamics) arranged in concentric rings. Backbone: standard causal transformer (token mixing = attention, tied LM head). Novelty: the per-layer MLP is replaced by a QuazimotoBlock that maps the hidden state to N oscillator phases on rings [4, 8, 16, 32, ...], runs a few differentiable Euler steps of structured Kuramoto coupling with a learnable frustration alpha, then reads out [cos, sin] of the phases. Set ring coupling to dense + alpha=0 and it degenerates toward AKOrN-style oscillatory neurons; the ring structure + hierarchy-to-center is the Quazimoto bit. """ from dataclasses import dataclass, field import math import os import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint _PKG_DIR = os.path.dirname(os.path.abspath(__file__)) # family traits (DERF soft_clamp, RMSNorm, and the opt-in refinement/aux blocks) from family import (soft_clamp, RMSNorm, HRMRefinementBlock, MoESwiGLU, MTPHead, JEPAPredictorBlock, PhaseAttentionRing, EngramRing, RingControllerBank, RingSpecialists) import instrument as _viz # live-visualizer capture hooks (no-op unless a recorder is active) from chimera import ChimeraBlock # grand-finale mixer: council of every project's physics @dataclass class QuazimotoConfig: vocab_size: int = 32000 n_layer: int = 10 n_head: int = 12 d_model: int = 768 block_size: int = 1024 dropout: float = 0.0 # Quazimoto oscillator bank ring_sizes: tuple = (4, 8, 16, 32, 64, 128, 256) # (unused by the Hydra tip mixer) osc_steps: int = 4 # differentiable Kuramoto Euler steps osc_dt: float = 0.5 readout_mult: int = 3 # MLP expansion on the readout # ---- Chimera mixer: a council of three physics cores, GRPO-selected (see chimera.py) ---- # KuramotoCore (Quazimoto oscillators) chim_osc: int = 48 # oscillators in the mean-field Kuramoto ring chim_osc_steps: int = 4 # Kuramoto Euler steps # GrowthCore (Mycel Neighbour-Sensing growth) chim_tips: int = 24 # tips growing in the bounded region chim_pos_dim: int = 2 # latent growth dimensionality chim_centers: int = 8 # low-rank density-field centres (O(N*F)) chim_growth_steps: int = 3 # growth Euler steps chim_growth_dt: float = 0.3 # growth step size # WaveCore (Wheeler-DeWitt) reuses wdw_* below. # GRPO council knobs (grpo_lm) chim_adv_clip: float = 2.0 # clamp on group-relative advantage chim_select_temp: float = 1.0 # selection softmax temperature chim_clip: float = 0.5 # PPO clip of selection weights around uniform chim_anchor: float = 0.25 # KL-to-uniform anchor strength # Wheeler-DeWitt wave core: K minisuperspace modes under a Lorentzian DeWitt supermetric. wdw_modes: int = 48 # K minisuperspace modes wdw_steps: int = 4 # leapfrog wave steps wdw_dt: float = 0.5 # base intrinsic-time step (scaled by the learnable lapse) wdw_constraint_weight: float = 0.05 # weight on the Hamiltonian-constraint aux loss # Mandelbrot phase seeding: seed the K minisuperspace modes from each token's fractal # orbit angles (frozen, parameter-free), added through a zero-init gate. use_fractal_phase_seed: bool = False # family-trait discipline osc_bound: float = 10.0 # DERF soft_clamp bound on block output (BPTT-stability) gate_init_open: float = 0.9 # tanh-gate opening at init (Quazimoto is the main mixer, # so it starts OPEN; set 0.0 for a pure no-op refinement) # interstitial rings: between each pair of oscillator rings, insert a phase-attention # ring + an engram ring that ABSORB info and INJECT it into the two neighbor rings. use_rings: bool = False ring_attn_heads: int = 4 ring_attn_head_dim: int = 16 ring_engram_compress: int = 32 ring_engram_heads: int = 2 ring_engram_table: int = 2048 ring_engram_ngram: int = 3 # per-ring memory specialists: a small MoE of `ring_n_specialists` mini experts # PER oscillator ring, each holding test-time-writable input/output stores that # act as an addressable context memory at inference (see family.RingSpecialists). use_ring_specialists: bool = False ring_n_specialists: int = 7 ring_spec_key_dim: int = 32 ring_spec_slot_dim: int = 64 ring_spec_top_k: int = 2 ring_spec_write_lr: float = 0.1 # per-ring self-organizing controllers (~1M total: one tiny net per oscillator # ring, shared across layers, weights self-optimize by a predictive/free-energy rule) use_ring_controllers: bool = False ring_ctrl_feat: int = 384 # fast-weight predictor width (dominates the ~150k/ring) ring_ctrl_local_lr: float = 0.01 # delta-rule (surprise-minimization) step size # opt-in family trait modules (all safe no-ops at init) use_hrm: bool = False hrm_steps: int = 3 hrm_dim: int = 256 hrm_gate_init: float = 0.1 # HRM gates start OPEN (random initial state needs a path out) use_moe: bool = False moe_intermediate: int = 768 moe_n_routed: int = 4 moe_n_shared: int = 1 moe_top_k: int = 2 use_mtp: bool = False mtp_layers: int = 4 # depth of MTP draft heads (predict +1..+mtp_layers); used for spec-decode mtp_loss_weight: float = 0.3 use_jepa: bool = False jepa_pred_dim: int = 256 jepa_horizon: int = 1 jepa_loss_weight: float = 0.1 # ---- attention: ported from model_v2.MLADerfXSAAttention ---- head_dim: int = 64 qk_rope_head_dim: int = 32 # partial RoPE: rope part of each head nope_head_dim: int = 32 # no-pos part (sum == head_dim) num_key_value_heads: int = 4 # GQA (== n_head for full MHA) q_lora_rank: int = 384 # MLA low-rank q projection o_lora_rank: int = 384 # MLA low-rank output projection use_qk_norm: bool = True # per-head RMSNorm on Q/K before RoPE use_derf: bool = False # erf attention instead of softmax (ablate) use_xsa: bool = False # value-subspace removal (ablate) # Elo attention (AlphaFold-style ranking / branch-until-winner): softmax attention # IS a Bradley-Terry tournament over keys; this gives each key a PERSISTENT rating # that accumulates across the sequence (winners of attention get boosted for later # queries) and feeds back as a logit bias. Rating is carried in the KV cache so the # tournament continues through incremental decoding. Zero-init gain -> no-op at start. use_elo: bool = False elo_k_init: float = 0.5 # Elo K-factor (rating step per "match"); per-head, learnable rope_theta: float = 10000.0 max_position_embeddings: int = 4096 rms_norm_eps: float = 1e-6 initializer_range: float = 0.02 # ---- model-level v2 features ---- zloss_coef: float = 1e-4 # z-loss on lm logits use_value_embed: bool = False # per-block value-embedding residual (zero-init gate) use_hyper_connections: bool = False hc_mult: int = 2 grad_checkpoint: bool = False # recompute each layer in backward (long-context / big-batch) @property def n_osc(self): # fractal seed must be wide enough for the widest core (each core slices what it needs) return max(self.chim_osc, self.chim_tips * self.chim_pos_dim, self.wdw_modes) def _ring_index(ring_sizes): """Return a LongTensor of length N mapping each oscillator -> its ring id.""" idx = [] for r, n in enumerate(ring_sizes): idx += [r] * n return torch.tensor(idx, dtype=torch.long) class QuazimotoBlock(nn.Module): """Oscillatory channel-mixing block: hidden state -> ring phases -> Kuramoto -> readout.""" def __init__(self, cfg: QuazimotoConfig): super().__init__() self.cfg = cfg N = cfg.n_osc R = len(cfg.ring_sizes) self.N, self.R = N, R self.norm = RMSNorm(cfg.d_model) # project hidden -> initial phases and natural frequencies self.to_theta = nn.Linear(cfg.d_model, N) self.to_omega = nn.Linear(cfg.d_model, N) # structured, LEARNABLE coupling: an RxR gain between ring blocks. # Block-constant coupling lets us use per-ring mean fields (order # parameters) -> O(N*R) instead of the O(N^2) pairwise sum. rid = _ring_index(cfg.ring_sizes) self.register_buffer("ring_id", rid) onehot = F.one_hot(rid, R).float() # [N,R] ring membership self.register_buffer("onehot", onehot) self.block_gain = nn.Parameter(self._init_gain(R)) self.center_phase = nn.Parameter(torch.zeros(())) # the rotating center ball self.center_omega = nn.Parameter(torch.tensor(0.3)) self.center_gain = nn.Parameter(torch.full((R,), 0.5)) # ring -> center pull self.alpha = nn.Parameter(torch.zeros(R)) # per-ring frustration # readout MLP on [cos theta, sin theta] hidden = cfg.readout_mult * cfg.d_model self.readout = nn.Sequential( nn.Linear(2 * N, hidden), nn.GELU(), nn.Linear(hidden, cfg.d_model), ) self.drop = nn.Dropout(cfg.dropout) # family-trait gate: out = readout * tanh(gate). gate_init_open seeds it so the # block is a warm-startable signal whose magnitude the optimizer controls; the # content weights are nonzero so the gate always receives gradient. go = math.atanh(min(cfg.gate_init_open, 0.9)) if cfg.gate_init_open > 0 else 0.0 self.gate = nn.Parameter(torch.full((1,), go)) # Mandelbrot phase-seed gate: theta <- to_theta(h) + tanh(seed_gate) * orbit. # Zero-init -> no-op at start; opens only if the fractal prior helps. if cfg.use_fractal_phase_seed: self.phase_seed_gate = nn.Parameter(torch.zeros(1)) # std=0.02 init discipline (matches the family's linear-layer init) for m in self.modules(): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) # interstitial rings (instantiated AFTER the std-init loop). Oscillators are # ordered ring-by-ring, so rings r and r+1 are the contiguous slice # [off_r, off_{r+2}); the slot between them injects into exactly that block. self.use_rings = cfg.use_rings self.slot_bounds = [] if self.use_rings: off = [0] for n in cfg.ring_sizes: off.append(off[-1] + n) self.slot_bounds = [(off[r], off[r + 2]) for r in range(R - 1)] self.attn_rings = nn.ModuleList([ PhaseAttentionRing(e - s, cfg.ring_attn_heads, cfg.ring_attn_head_dim, cfg.osc_bound) for (s, e) in self.slot_bounds]) self.engram_rings = nn.ModuleList([ EngramRing(cfg.d_model, e - s, cfg.ring_engram_compress, cfg.ring_engram_heads, cfg.ring_engram_table, cfg.ring_engram_ngram, cfg.osc_bound) for (s, e) in self.slot_bounds]) # one memory-specialist MoE bank PER oscillator ring (injects into that ring) self.use_ring_specialists = cfg.use_ring_specialists if self.use_ring_specialists: self.specialists = nn.ModuleList([ RingSpecialists(n, cfg.d_model, cfg.ring_n_specialists, cfg.ring_spec_key_dim, cfg.ring_spec_slot_dim, cfg.ring_spec_top_k, cfg.ring_spec_write_lr, cfg.osc_bound) for n in cfg.ring_sizes]) @staticmethod def _init_gain(R): # within-ring coupling strong on the diagonal, weak neighbor coupling g = torch.zeros(R, R) for i in range(R): g[i, i] = 1.0 if i + 1 < R: g[i, i + 1] = g[i + 1, i] = 0.3 return g def forward(self, x, ring_ctl=None, phase_seed=None): cfg = self.cfg B, T, _ = x.shape h = self.norm(x) theta = self.to_theta(h) # [B,T,N] initial phases if phase_seed is not None: # Mandelbrot orbit phase prior (gated, no-op at init) theta = theta + torch.tanh(self.phase_seed_gate) * phase_seed omega = torch.tanh(self.to_omega(h)) # [B,T,N] natural frequencies (bounded) # per-ring self-organizing controller: observe each ring (order-parameter # magnitude/phase + mean freq), get no-op-at-init modulations of that ring's # coupling / frustration / center-pull / injection. d_coup = d_alpha = d_center = d_inj = None if ring_ctl is not None: with torch.no_grad(): c0, s0 = torch.cos(theta), torch.sin(theta) cnt = self.onehot.sum(0).clamp(min=1.0) # [R] ring sizes zc = (c0 @ self.onehot) / cnt # [B,T,R] zs = (s0 @ self.onehot) / cnt mag = torch.sqrt(zc ** 2 + zs ** 2 + 1e-8) om = (omega @ self.onehot) / cnt obs = torch.stack([mag.mean((0, 1)), (zc / mag).mean((0, 1)), (zs / mag).mean((0, 1)), om.mean((0, 1))], dim=-1) # [R,4] ctrl = ring_ctl(obs) # [R,4] (grad -> enc/dec only) d_coup = torch.tanh(ctrl[:, 0]) # ring coupling scale (1 + .) d_alpha = 0.5 * torch.tanh(ctrl[:, 1]) # frustration shift d_center = torch.tanh(ctrl[:, 2]) # center-pull scale (1 + .) d_inj = torch.tanh(ctrl[:, 3]) # injection scale (1 + .) # interstitial rings absorb info ONCE (from the initial phases / hidden) and # inject a constant drive into their two neighbor oscillator rings' phase update. inject = 0.0 if self.use_rings: inject = torch.zeros_like(theta) for (s, e), ar, er in zip(self.slot_bounds, self.attn_rings, self.engram_rings): contrib = ar(theta[..., s:e]) + er(h) # [B,T,e-s] inject = inject + F.pad(contrib, (s, self.N - e)) # place at [s:e], sum overlaps if d_inj is not None: inject = inject * (1.0 + d_inj)[self.ring_id] # controller scales the drive # memory specialists: each per-ring bank injects retrieved store_out into its # ring's oscillators. Rings are contiguous and tile [0,N), so cat == full drive. if self.use_ring_specialists: spec_inj = torch.cat([bank(h) for bank in self.specialists], dim=-1) # [B,T,N] inject = inject + spec_inj rid = self.ring_id # [N] gain = self.block_gain if d_coup is None else self.block_gain * (1.0 + d_coup).unsqueeze(1) G = gain[rid] # [N,R] per-oscillator ring gains alpha_vec = self.alpha if d_alpha is None else self.alpha + d_alpha alpha_i = alpha_vec[rid] # [N] center = self.center_phase # scalar phase, advances in time cpull = self.center_gain if d_center is None else self.center_gain * (1.0 + d_center) center_pull = cpull[rid] # [N] invN = 1.0 / self.N for s in range(cfg.osc_steps): ctr = center + self.center_omega * (s * cfg.osc_dt) c, sn = torch.cos(theta), torch.sin(theta) # [B,T,N] # per-ring summed mean fields, then routed back through gains G zc = (c @ self.onehot) @ G.t() # [B,T,N] zs = (sn @ self.onehot) @ G.t() # [B,T,N] A = alpha_i - theta # [B,T,N] # (1/N) * sum_r G_ir * Im( e^{i(alpha_i - theta_i)} * Zsum_r ) coupling = invN * (torch.sin(A) * zc + torch.cos(A) * zs) to_center = center_pull * torch.sin(ctr - theta + alpha_i) # [B,T,N] theta = theta + cfg.osc_dt * (omega + coupling + to_center + inject) # ring drive feat = torch.cat([torch.cos(theta), torch.sin(theta)], dim=-1) # [B,T,2N] out = self.readout(feat) * torch.tanh(self.gate) out = self.drop(soft_clamp(out, self.cfg.osc_bound)) # bounded, live grad rec = _viz.get_rec() if rec is not None and rec.enabled: # live-viz capture th = theta[0, -1] # last-token phases [N] cnt = self.onehot.sum(0).clamp(min=1.0) zc = (torch.cos(th) @ self.onehot) / cnt # per-ring order param zs = (torch.sin(th) @ self.onehot) / cnt R = torch.sqrt(zc ** 2 + zs ** 2 + 1e-9) psi = torch.atan2(zs, zc) rec.log_ring(R.tolist(), psi.tolist(), th.tolist()) rec.log_quaz_norm(out[0, -1].norm().item()) rec.flush_spec() # group this layer's routes return out class RotaryEmbedding(nn.Module): """RoPE for the rope partition of Q/K (qk_rope_head_dim dims). Ported from v2.""" def __init__(self, dim, max_positions=4096, theta=10000.0): super().__init__() inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) t = torch.arange(max_positions).float() freqs = torch.outer(t, inv_freq) self.register_buffer("cos_cache", freqs.cos(), persistent=False) self.register_buffer("sin_cache", freqs.sin(), persistent=False) def forward(self, x, position_ids): cos = self.cos_cache[position_ids].unsqueeze(1) sin = self.sin_cache[position_ids].unsqueeze(1) d = cos.shape[-1] x1, x2 = x[..., :d], x[..., d:] return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) class MLADerfXSAAttention(nn.Module): """Family attention, ported from model_v2.MLADerfXSAAttention: MLA low-rank q/o projections, partial RoPE (nope+rope split), QK-Norm, optional DERF erf attention, optional XSA value-subspace removal, GQA. Supports an optional KV cache for incremental decoding (cache holds post-RoPE, pre-GQA-expansion k/v of shape [B, num_kv_heads, S, head_dim]).""" def __init__(self, cfg: QuazimotoConfig): super().__init__() self.num_heads = cfg.n_head self.num_kv_heads = cfg.num_key_value_heads self.head_dim = cfg.head_dim self.nope_head_dim = cfg.nope_head_dim self.use_qk_norm = cfg.use_qk_norm self.use_derf = cfg.use_derf self.use_xsa = cfg.use_xsa self.use_elo = cfg.use_elo self.dropout_p = cfg.dropout self.kv_groups = self.num_heads // self.num_kv_heads assert self.nope_head_dim + cfg.qk_rope_head_dim == self.head_dim, \ "nope_head_dim + qk_rope_head_dim must equal head_dim" self.q_a_proj = nn.Linear(cfg.d_model, cfg.q_lora_rank, bias=False) self.q_a_norm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps) self.q_b_proj = nn.Linear(cfg.q_lora_rank, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(cfg.d_model, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(cfg.d_model, self.num_kv_heads * self.head_dim, bias=False) self.o_a_proj = nn.Linear(self.num_heads * self.head_dim, cfg.o_lora_rank, bias=False) self.o_b_proj = nn.Linear(cfg.o_lora_rank, cfg.d_model, bias=False) if self.use_qk_norm: self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) self.rope = RotaryEmbedding(cfg.qk_rope_head_dim, cfg.max_position_embeddings, cfg.rope_theta) if self.use_derf: self.derf_alpha = nn.Parameter(torch.ones(self.num_heads)) self.derf_bias = nn.Parameter(torch.zeros(self.num_heads)) self.derf_gamma = nn.Parameter(torch.ones(self.num_heads)) if self.use_elo: # per-head Elo K-factor (softplus>0) and rating->logit gain (tanh, zero-init # => no-op at start, so the tournament reduces to plain softmax attention). self.elo_log_k = nn.Parameter(torch.full((self.num_heads,), math.log(math.expm1(cfg.elo_k_init)))) self.elo_gain = nn.Parameter(torch.zeros(self.num_heads)) for m in (self.q_a_proj, self.q_b_proj, self.k_proj, self.v_proj, self.o_a_proj, self.o_b_proj): nn.init.normal_(m.weight, std=cfg.initializer_range) def _elo_attention(self, q, k, v, offset, S, L, past_rating, device): """Elo/Bradley-Terry tournament attention. Softmax attention already gives each key its win-probability against the field; Elo adds a PERSISTENT rating r_j that accumulates across queries and biases future logits (winners get boosted), exactly like the AlphaFold-style branch-until-winner reinforcement. e_ij = q_i.k_j/sqrt(d) + gain * rating_j (rating = prior + intra-call) match : S_ij = softmax weight key j won at query i ; E_i = 1/n_i (uniform baseline) update : dR_ij = K * (S_ij - E_i) (over-performers gain rating, under-perform lose) rating_j(query i) = sum over earlier queries t pass B == pass A == plain softmax (family no-op contract). Returns (y, rating_out) with rating_out: [B, num_heads, L].""" B, H = q.shape[0], self.num_heads scale = 1.0 / math.sqrt(self.head_dim) scores = torch.matmul(q, k.transpose(-2, -1)) * scale # [B,H,S,L] qpos = torch.arange(S, device=device).view(S, 1) + offset kpos = torch.arange(L, device=device).view(1, L) keep = (kpos <= qpos) # [S,L] causal keepf = keep.to(scores.dtype) neg = torch.finfo(scores.dtype).min K = F.softplus(self.elo_log_k).view(1, H, 1, 1) # per-head K-factor >0 gain = torch.tanh(self.elo_gain).view(1, H, 1, 1) # per-head, 0 at init # prior rating of each key = what it accumulated BEFORE this forward (from the # cache) + 0 for the keys introduced this call. [B,H,L] -> bias [B,H,1,L]. if past_rating is not None: prior = torch.cat([past_rating, past_rating.new_zeros(B, H, S)], dim=-1) # new keys start at 0 else: prior = scores.new_zeros(B, H, L) prior_bias = gain * prior.unsqueeze(2) # [B,H,1,L] masked = scores.masked_fill(~keep, neg) # Match outcome = innate q.k compatibility (the "game result"), NOT the accumulated # rating. Keeping matches rating-INDEPENDENT is what makes the parallel prefix-sum # (tril @ dR, training) bit-identical to the incremental running sum (decoding); # rating is the REPUTATION that biases attention, not what defines the match. a_base = torch.softmax(masked, dim=-1) # [B,H,S,L] n_i = keepf.sum(-1).clamp(min=1.0).view(1, 1, S, 1) # valid keys per query dR = (K * (a_base - 1.0 / n_i)) * keepf # over/under-perform vs field # rating seen by query i = prior (from cache) + causal sum of dR from earlier queries t 0 and self.training: a = F.dropout(a, p=self.dropout_p) y = torch.matmul(a, v) # [B,H,S,hd] # rating carried forward = prior + everything this call's queries contributed rating_out = prior + dR.sum(2) # [B,H,L] return y, rating_out def forward(self, x, position_ids, past_kv=None, use_cache=False): B, S, _ = x.shape q = self.q_b_proj(self.q_a_norm(self.q_a_proj(x))) q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2) if self.use_qk_norm: q, k = self.q_norm(q), self.k_norm(k) d = self.nope_head_dim q = torch.cat([q[..., :d], self.rope(q[..., d:], position_ids)], dim=-1) k = torch.cat([k[..., :d], self.rope(k[..., d:], position_ids)], dim=-1) # KV cache: prepend previously-seen post-RoPE k/v (pre-GQA-expansion so # the cache stays GQA-compact), then this step's tokens become the tail. past_rating = None if past_kv is not None: if len(past_kv) == 3: # (k, v, elo_rating) cache past_k, past_v, past_rating = past_kv else: past_k, past_v = past_kv k = torch.cat([past_k, k], dim=2) v = torch.cat([past_v, v], dim=2) present = (k, v) if use_cache else None L = k.size(2) # total keys (cache + current) if self.kv_groups > 1: k = k.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape( B, self.num_heads, L, self.head_dim) v = v.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape( B, self.num_heads, L, self.head_dim) # Causal mask over the [S queries x L keys] block. When there is no cache # and S == L this is the plain lower triangle; with a cache the S new # queries sit at absolute positions [L-S, L) and may attend all keys <= # their own position. allowed[i,j] = j <= (L - S + i). offset = L - S if self.use_elo: # Elo tournament over keys y, rating_out = self._elo_attention(q, k, v, offset, S, L, past_rating, x.device) if use_cache: present = (present[0], present[1], rating_out) elif self.use_derf: qpos = torch.arange(S, device=x.device).view(S, 1) + offset kpos = torch.arange(L, device=x.device).view(1, L) is_masked = (kpos > qpos).unsqueeze(0).unsqueeze(0) # [1,1,S,L] scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) safe = scores.masked_fill(is_masked, -10000.0) a, b, g = (self.derf_alpha.view(1, -1, 1, 1), self.derf_bias.view(1, -1, 1, 1), self.derf_gamma.view(1, -1, 1, 1)) w = g * torch.erf(a * safe + b) w = (w + g) / 2.0 w = w.masked_fill(is_masked, 0.0) w = w / (w.sum(-1, keepdim=True) + 1e-8) if self.dropout_p > 0 and self.training: w = F.dropout(w, p=self.dropout_p) y = torch.matmul(w, v) else: if offset == 0: attn_mask, causal = None, True else: qpos = torch.arange(S, device=x.device).view(S, 1) + offset kpos = torch.arange(L, device=x.device).view(1, L) attn_mask = (kpos <= qpos).unsqueeze(0).unsqueeze(0) # bool: True=keep causal = False y = F.scaled_dot_product_attention( q.contiguous(), k.contiguous(), v.contiguous(), attn_mask=attn_mask, is_causal=causal, dropout_p=self.dropout_p if self.training else 0.0) rec = _viz.get_rec() if rec is not None and rec.enabled: # last-query attention, mean over heads sc = (q[:, :, -1:] @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) # [B,H,1,L] w = torch.softmax(sc, dim=-1).mean(1)[0, 0] # [L] rec.log_attn(getattr(self, "layer_idx", 0), w.tolist()) if self.use_xsa: # remove each query's component along its OWN value direction; with a # cache the queries are the last S of the L cached value positions. vq = v[:, :, -S:, :] vn = vq / (vq.norm(dim=-1, keepdim=True) + 1e-8) y = y - (y * vn).sum(-1, keepdim=True) * vn y = y.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim) out = self.o_b_proj(self.o_a_proj(y)) return (out, present) if use_cache else out class HyperConnectionLayer(nn.Module): """Softmax pre-mix / post-distribute over hc_mult residual streams (v2).""" def __init__(self, hc_mult): super().__init__() self.pre_weight = nn.Parameter(torch.linspace(0.5, -0.5, hc_mult) / max(hc_mult, 1)) self.post_weight = nn.Parameter(torch.linspace(-0.5, 0.5, hc_mult) / max(hc_mult, 1)) def pre_op(self, copies): w = F.softmax(self.pre_weight, dim=0) return (copies * w.view(1, -1, 1, 1)).sum(dim=1) def post_op(self, copies, delta): w = F.softmax(self.post_weight, dim=0) return copies + delta.unsqueeze(1) * w.view(1, -1, 1, 1) class HCOutputMix(nn.Module): """Learned softmax mix over hc_mult streams at the output (init == mean).""" def __init__(self, hc_mult): super().__init__() self.weight = nn.Parameter(torch.zeros(hc_mult)) def forward(self, copies): w = F.softmax(self.weight, dim=0) return (copies * w.view(1, -1, 1, 1)).sum(dim=1) class Layer(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.use_hc = cfg.use_hyper_connections self.use_value_embed = cfg.use_value_embed self.attn_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps) self.attn = MLADerfXSAAttention(cfg) self.quaz = ChimeraBlock(cfg) # the Chimera council mixer sub-layer if self.use_hc: self.hc_attn = HyperConnectionLayer(cfg.hc_mult) self.hc_ffn = HyperConnectionLayer(cfg.hc_mult) if self.use_value_embed: self.ve_gate = nn.Parameter(torch.zeros(1)) # zero-init -> no-op def forward(self, x, position_ids, token_embed=None, ring_ctl=None, past_kv=None, use_cache=False, phase_seed=None): h = self.hc_attn.pre_op(x) if self.use_hc else x if self.use_value_embed and token_embed is not None: h = h + torch.tanh(self.ve_gate) * token_embed attn_out = self.attn(self.attn_norm(h), position_ids, past_kv=past_kv, use_cache=use_cache) present = None if use_cache: attn_out, present = attn_out if self.use_hc: x = self.hc_attn.post_op(x, attn_out) h = self.hc_ffn.pre_op(x) else: h = h + attn_out ffn_out = self.quaz(h, ring_ctl, phase_seed) # QuazimotoBlock norms internally if self.use_hc: x = self.hc_ffn.post_op(x, ffn_out) else: x = h + ffn_out return (x, present) if use_cache else x class QuazimotoLM(nn.Module): def __init__(self, cfg: QuazimotoConfig): super().__init__() self.cfg = cfg self.tok = nn.Embedding(cfg.vocab_size, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) # positions come from RoPE (no abs embed) # frozen Mandelbrot orbit-angle phase table [vocab, n_osc]; non-persistent # (deterministic, regenerated on load -> not stored in checkpoints). if cfg.use_fractal_phase_seed: from fractal import mandelbrot_phase_table, load_phase_table pt, mode = load_phase_table(cfg.vocab_size, cfg.n_osc, os.path.join(_PKG_DIR, "fractal_phase.pt")) if pt is None: pt = mandelbrot_phase_table(cfg.vocab_size, cfg.n_osc) print(" [fractal] flat Halton seed (run build_fractal_table.py for hierarchical)") else: print(f" [fractal] loaded {mode} phase table") self.register_buffer("fractal_phase", pt, persistent=False) self.layers = nn.ModuleList([Layer(cfg) for _ in range(cfg.n_layer)]) for i, layer in enumerate(self.layers): layer.attn.layer_idx = i # for live-viz attention capture self.norm_f = RMSNorm(cfg.d_model, cfg.rms_norm_eps) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.head.weight = self.tok.weight # weight tying self.logit_scale = nn.Parameter(torch.tensor(1.0)) # family trait: learned logit temp self.use_value_embed = cfg.use_value_embed self.hc_out_mix = HCOutputMix(cfg.hc_mult) if cfg.use_hyper_connections else None self.apply(self._init) # engram rings manage their own inits (frozen compressor, std-0.01 tables, # DERF bias -4); re-apply them after the global init clobbers those. for m in self.modules(): if isinstance(m, (EngramRing, RingSpecialists)): m.family_reinit() # opt-in trait modules -- instantiated AFTER self.apply so their internal # zero-inits (MoE down-proj, MTP proj) and gate inits survive the global init. # HRM/MoE refine the trunk after the stack (the family's "top membrane # refinement" point); MTP/JEPA are train-time aux heads on the trunk rep. self.hrm = (HRMRefinementBlock(cfg.d_model, cfg.hrm_dim, cfg.hrm_steps, gate_init_open=cfg.hrm_gate_init) if cfg.use_hrm else None) self.moe = MoESwiGLU(cfg.d_model, cfg.moe_intermediate, cfg.moe_n_routed, cfg.moe_n_shared, cfg.moe_top_k) if cfg.use_moe else None self.mtp_heads = (nn.ModuleList([MTPHead(cfg.d_model) for _ in range(cfg.mtp_layers)]) if cfg.use_mtp else None) self.jepa = (JEPAPredictorBlock(cfg.d_model, cfg.jepa_pred_dim, cfg.jepa_horizon) if cfg.use_jepa else None) self.ring_bank = (RingControllerBank(len(cfg.ring_sizes), 4, cfg.ring_ctrl_feat, cfg.ring_ctrl_local_lr) if cfg.use_ring_controllers else None) # exact kwargs to rebuild on resume (family `family_config` convention) self.family_config = {k: getattr(cfg, k) for k in ( "vocab_size", "n_layer", "n_head", "d_model", "block_size", "ring_sizes", "osc_steps", "osc_dt", "readout_mult", "osc_bound", "gate_init_open", "use_hrm", "hrm_steps", "hrm_dim", "hrm_gate_init", "use_moe", "moe_intermediate", "moe_n_routed", "moe_n_shared", "moe_top_k", "use_mtp", "mtp_layers", "mtp_loss_weight", "use_jepa", "jepa_pred_dim", "jepa_horizon", "jepa_loss_weight", "head_dim", "qk_rope_head_dim", "nope_head_dim", "num_key_value_heads", "q_lora_rank", "o_lora_rank", "use_qk_norm", "use_derf", "use_xsa", "use_elo", "elo_k_init", "rope_theta", "max_position_embeddings", "rms_norm_eps", "initializer_range", "zloss_coef", "use_value_embed", "use_hyper_connections", "hc_mult", "use_rings", "ring_attn_heads", "ring_attn_head_dim", "ring_engram_compress", "ring_engram_heads", "ring_engram_table", "ring_engram_ngram", "use_ring_controllers", "ring_ctrl_feat", "ring_ctrl_local_lr", "use_ring_specialists", "ring_n_specialists", "ring_spec_key_dim", "ring_spec_slot_dim", "ring_spec_top_k", "ring_spec_write_lr", "use_fractal_phase_seed", # Chimera council geometry -- so checkpoints rebuild the exact three cores "chim_osc", "chim_osc_steps", "chim_tips", "chim_pos_dim", "chim_centers", "chim_growth_steps", "chim_growth_dt", "chim_adv_clip", "chim_select_temp", "chim_clip", "chim_anchor", "wdw_modes", "wdw_steps", "wdw_dt", "wdw_constraint_weight")} n = sum(p.numel() for p in self.parameters()) print(f"Chimera-LM: {n/1e6:.1f}M params | council mixer: Kuramoto({cfg.chim_osc}) + " f"Growth({cfg.chim_tips}t) + Wave({cfg.wdw_modes}) GRPO-selected per token | " f"traits: {self.active_traits()}") def active_traits(self): t = [] if self.cfg.use_hrm: t.append("HRM") if self.cfg.use_moe: t.append("MoE") if self.cfg.use_mtp: t.append("MTP") if self.cfg.use_jepa: t.append("JEPA") if self.cfg.use_rings: t.append("Rings") if self.cfg.use_ring_controllers: t.append("Controllers") if self.cfg.use_ring_specialists: t.append(f"Specialists({self.cfg.ring_n_specialists}/ring)") if self.cfg.use_fractal_phase_seed: t.append("FractalSeed") return ",".join(t) or "none" def reset_ring_memory(self): """Zero every RingSpecialists store. Call between independent prompts so the test-time context memory does not carry over from a previous sequence.""" for m in self.modules(): if isinstance(m, RingSpecialists): m.reset_memory() def set_ring_memory_writing(self, enabled: bool): """Toggle online writes to the specialist stores (e.g. freeze the memory during evaluation, or disable it to ablate the test-time-write behavior).""" for m in self.modules(): if isinstance(m, RingSpecialists): m.write_enabled = enabled def _init(self, m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=0.02) def _compute_trunk(self, idx, past_key_values=None, use_cache=False): """Shared backbone: tokens -> layers -> trunk refinements -> RMSNorm trunk. Returns (trunk, presents). Used by both forward() and forward_drafts().""" B, T = idx.shape # with a cache, the new tokens sit at absolute positions [past_len, past_len+T) past_len = past_key_values[0][0].size(2) if past_key_values else 0 position_ids = (torch.arange(past_len, past_len + T, device=idx.device) .unsqueeze(0).expand(B, -1)) x = self.drop(self.tok(idx)) token_embed = x if self.use_value_embed else None phase_seed = self.fractal_phase[idx] if self.cfg.use_fractal_phase_seed else None if self.cfg.use_hyper_connections: x = x.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1, -1).clone() # [B,M,T,H] presents = [] if use_cache else None ckpt = getattr(self.cfg, "grad_checkpoint", False) and self.training and not use_cache for i, layer in enumerate(self.layers): past = past_key_values[i] if past_key_values is not None else None if ckpt: # recompute each layer in backward -> only ONE layer's activations (incl. # the Elo O(T^2) score matrix) are live at a time. Enables long-context / # big-batch training that would otherwise OOM. No effect on cached decode. def _run(inp, _layer=layer): return _layer(inp, position_ids, token_embed, self.ring_bank, past_kv=None, use_cache=False, phase_seed=phase_seed) x = torch.utils.checkpoint.checkpoint(_run, x, use_reentrant=False) else: x = layer(x, position_ids, token_embed, self.ring_bank, past_kv=past, use_cache=use_cache, phase_seed=phase_seed) if use_cache: x, present = x presents.append(present) if self.hc_out_mix is not None: x = self.hc_out_mix(x) # learned mix over streams (init == mean) # trunk refinements (no-op at init): HRM iterative gated refine, MoE expert mix rec = _viz.get_rec() if self.hrm is not None: hx = self.hrm(x) if rec is not None and rec.enabled: rec.log_trait("hrm", (hx - x)[0, -1].norm().item()) x = hx if self.moe is not None: mx = self.moe(x) if rec is not None and rec.enabled: rec.log_trait("moe", mx[0, -1].norm().item()) x = x + mx return self.norm_f(x), presents # v2: RMSNorm trunk (no tanh) @torch.no_grad() def forward_drafts(self, idx, past_key_values=None, use_cache=False): """Speculative-decoding support (DeepSpec-style draft+verify, self-drafted): return (main_logits, mtp_logits_list[, presents]) over ALL T positions. * main_logits[:, j] predicts the token at position j+1 (verifier) * mtp_logits_list[k][:, j] predicts the token at position j+2+k (k=0..mtp-1) So one forward yields, at the last position, the genuine next token (main) plus `mtp_layers` drafted future tokens to be verified next round.""" trunk, presents = self._compute_trunk(idx, past_key_values, use_cache) main_logits = self.head(trunk) * self.logit_scale mtp_logits = ([self.head(h(trunk)) * self.logit_scale for h in self.mtp_heads] if self.mtp_heads is not None else []) return (main_logits, mtp_logits, presents) if use_cache else (main_logits, mtp_logits) def forward(self, idx, targets=None, past_key_values=None, use_cache=False): B, T = idx.shape trunk, presents = self._compute_trunk(idx, past_key_values, use_cache) logits = self.head(trunk) * self.logit_scale loss, aux = None, {} if targets is not None: flat = logits.view(-1, logits.size(-1)) tflat = targets.reshape(-1) loss = F.cross_entropy(flat, tflat, ignore_index=-1) # z-loss (v2): penalise log^2 of the partition function -> no logit drift if self.cfg.zloss_coef > 0: valid = tflat != -1 if valid.any(): log_z = torch.logsumexp(flat[valid].float(), dim=-1) aux["zloss"] = self.cfg.zloss_coef * (log_z ** 2).mean() if self.moe is not None and self.moe.last_aux_loss is not None: aux["moe_aux"] = self.moe.last_aux_loss # Wheeler-DeWitt Hamiltonian constraint: pressure every layer's block toward # the physical surface H = 0 (H Psi = 0). Sum over layers. if self.cfg.wdw_constraint_weight > 0: hsum, nh = 0.0, 0 for layer in self.layers: lc = getattr(layer.quaz, "last_constraint", None) if lc is not None: hsum = hsum + lc; nh += 1 if nh: aux["wdw_constraint"] = self.cfg.wdw_constraint_weight * hsum / nh if self.mtp_heads is not None: mtp_total, n_active = 0.0, 0 for k, head in enumerate(self.mtp_heads, start=1): if T - k <= 0: break mlog = self.head(head(trunk[:, :T - k])) * self.logit_scale mtp_total = mtp_total + F.cross_entropy( mlog.reshape(-1, self.cfg.vocab_size), targets[:, k:].reshape(-1), ignore_index=-1) n_active += 1 if n_active: aux["mtp_loss"] = self.cfg.mtp_loss_weight * mtp_total / n_active if self.jepa is not None and T > 1: jepa_total, n_j = 0.0, 0 for k in range(1, self.cfg.jepa_horizon + 1): if T - k <= 0: break pred = self.jepa(trunk[:, :T - k], k) tgt = trunk[:, k:].detach() # JEPA stop-grad target cos = F.cosine_similarity(pred.float(), tgt.float(), dim=-1) jepa_total = jepa_total + (1.0 - cos).mean() n_j += 1 if n_j: aux["jepa_loss"] = self.cfg.jepa_loss_weight * jepa_total / n_j # use_cache returns the per-layer (k, v) presents as a 4th item; without # it the signature is the original 3-tuple so train.py is unaffected. if use_cache: return logits, loss, aux, presents return logits, loss, aux @torch.no_grad() def generate(self, idx, n_new, temperature=1.0, top_k=None, use_cache=True): """Autoregressive sampling with the family's NaN/inf sanitisation: long accumulation can push a logit to +/-inf or NaN, whose softmax -> NaN -> multinomial fires a CUDA device-side assert. Clamp first, greedy-fallback. With use_cache=True (default) the attention KV cache is kept across steps so each new token costs one forward over a single position instead of a full recompute. The cache holds absolute positions, so once it fills past block_size we drop the cache and fall back to a windowed recompute (RoPE positions would otherwise exceed max_position_embeddings).""" self.eval() past = None for _ in range(n_new): if use_cache and past is not None and past[0][0].size(2) < self.cfg.block_size: step_in = idx[:, -1:] # only the newest token logits, _, _, past = self(step_in, past_key_values=past, use_cache=True) elif use_cache: cond = idx[:, -self.cfg.block_size:] # prefill / re-prime cache logits, _, _, past = self(cond, use_cache=True) else: cond = idx[:, -self.cfg.block_size:] logits, _, _ = self(cond) lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0, posinf=1e4, neginf=-1e4) / max(temperature, 1e-6) if top_k: k = min(top_k, lg.size(-1)) v, _ = torch.topk(lg, k) lg[lg < v[:, [-1]]] = -float("inf") probs = torch.softmax(lg, dim=-1) if not torch.isfinite(probs).all() or float(probs.sum()) <= 0.0: nxt = torch.argmax(lg, dim=-1, keepdim=True) else: nxt = torch.multinomial(probs, 1) idx = torch.cat([idx, nxt], dim=1) return idx if __name__ == "__main__": # exercise the base model and all trait modules at once cfg = QuazimotoConfig(use_hrm=True, use_moe=True, use_mtp=True, use_jepa=True) model = QuazimotoLM(cfg) x = torch.randint(0, cfg.vocab_size, (2, 64)) logits, loss, aux = model(x, x) total = loss + sum(aux.values()) print("forward ok:", tuple(logits.shape), "loss", round(loss.item(), 3), "| aux:", {k: round(float(v), 4) for k, v in aux.items()}) total.backward() print("backward ok")