""" Cozet: Native SYNAXIM Base Model (Enhanced) Architecture: SymbioticGate (M-matrix) + NaN Firewall + Dual-Track M (c) 2026 GRRN Research. """ import torch, torch.nn as nn, torch.nn.functional as F, math from dataclasses import dataclass @dataclass class CozConfig: hidden_size: int = 1024 num_layers: int = 12 num_attention_heads: int = 16 num_kv_heads: int = 4 intermediate_size: int = 4096 vocab_size: int = 50257 max_seq_len: int = 4096 rope_theta: float = 10000.0 rms_norm_eps: float = 1e-6 memory_decay: float = 0.995 unrotated_decay: float = 0.995 tie_word_embeddings: bool = True @property def head_dim(self): return self.hidden_size // self.num_attention_heads class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): return (x.float() * x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()).type_as(x) * self.weight class NaNFirewall(nn.Module): def __init__(self, clamp_value=1e4, norm_ratio=8.0): super().__init__() self.clamp_value = clamp_value self.norm_ratio = norm_ratio def forward(self, delta, h_ref): delta = torch.nan_to_num(delta, nan=0.0, posinf=self.clamp_value, neginf=-self.clamp_value) d_norm = delta.norm() max_norm = self.norm_ratio * (h_ref.norm() + 1e-8) if d_norm > max_norm: delta = delta * (max_norm / (d_norm + 1e-8)) if torch.isnan(delta).all() or torch.isinf(delta).all(): delta = torch.zeros_like(delta) return delta class SymbioticGate(nn.Module): def __init__(self, config, layer_idx): super().__init__() D, nh, nk, hd = config.hidden_size, config.num_attention_heads, config.num_kv_heads, config.head_dim self.D, self.n_heads, self.n_kv, self.head_dim = D, nh, nk, hd self.decay = config.memory_decay self.unrot_decay = config.unrotated_decay self.q_proj = nn.Linear(D, nh * hd, bias=False) self.k_proj = nn.Linear(D, nk * hd, bias=False) self.v_proj = nn.Linear(D, nk * hd, bias=False) self.o_proj = nn.Linear(nh * hd, D, bias=False) self.gate_scale = nn.Parameter(torch.ones(1)) self.gate_bias = nn.Parameter(torch.zeros(1)) self._rc, self._rs = None, None def _rope(self, max_pos, dev): if self._rc is not None and max_pos <= self._rc.shape[0]: return half = self.head_dim // 2 f = 1.0 / (10000.0 ** (torch.arange(0, half, device=dev).float() / half)) a = torch.outer(torch.arange(max_pos, device=dev).float(), f) self._rc, self._rs = a.cos(), a.sin() def _apply_rope(self, x, pos): half = self.head_dim // 2 c, s = self._rc[pos, :half], self._rs[pos, :half] x1, x2 = x[..., :half], x[..., half:] return torch.cat([x1*c - x2*s, x1*s + x2*c], dim=-1) def forward(self, h, M, M_unrot, pos): self._rope(pos + 1, h.device) q, k, v = self.q_proj(h), self.k_proj(h), self.v_proj(h) q_h = self._apply_rope(q.view(self.n_heads, self.head_dim), pos) k_h = self._apply_rope(k.view(self.n_kv, self.head_dim), pos) v_h = v.view(self.n_kv, self.head_dim) if self.n_kv < self.n_heads: r = self.n_heads // self.n_kv k_h, v_h = k_h.repeat_interleave(r, 0), v_h.repeat_interleave(r, 0) g = torch.sigmoid((q_h * k_h).sum(-1).mean() / math.sqrt(self.head_dim) * self.gate_scale + self.gate_bias) kf, vf = k_h.reshape(-1), v_h.reshape(-1) kn = kf / (kf.norm() + 1e-8) vn = vf / (vf.norm() + 1e-8) * h.norm() outer = torch.outer(kn, vn) M2 = g * self.decay * M + (1.0 - g) * outer M_unrot2 = self.unrot_decay * M_unrot + (1.0 - self.unrot_decay) * outer return self.o_proj(q_h.reshape(-1) @ M2), M2, M_unrot2 class SynaxBlock(nn.Module): def __init__(self, config, i): super().__init__() self.norm_attn = RMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = SymbioticGate(config, i) self.norm_mlp = RMSNorm(config.hidden_size, config.rms_norm_eps) self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) self.firewall = NaNFirewall() def forward(self, h, M, M_unrot, pos): a, M2, M_unrot2 = self.attn(self.norm_attn(h), M, M_unrot, pos) a = self.firewall(a, h) h = h + a n = self.norm_mlp(h) m = self.down_proj(F.silu(self.gate_proj(n)) * self.up_proj(n)) m = self.firewall(m, h) h = h + m return h, M2, M_unrot2 class CozModel(nn.Module): def __init__(self, config): super().__init__() self.config, self.D, self.n_layers = config, config.hidden_size, config.num_layers self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList([SynaxBlock(config, i) for i in range(config.num_layers)]) self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) if not config.tie_word_embeddings: self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) ds = 1.0 / math.sqrt(2 * self.n_layers) for name, p in self.named_parameters(): if p.dim() < 2: continue if "embed" in name: nn.init.normal_(p, std=0.02) elif "down_proj" in name or "o_proj" in name: nn.init.normal_(p, std=0.02 * ds) elif p.dim() == 2: nn.init.normal_(p, std=0.02) def init_m(self, dev): M = [torch.zeros(self.D, self.D, device=dev) for _ in range(self.n_layers)] M_unrot = [torch.zeros(self.D, self.D, device=dev) for _ in range(self.n_layers)] return M, M_unrot def forward_token(self, tid, M, M_unrot, pos): h = self.embed_tokens.weight[tid] for i, layer in enumerate(self.layers): h, M[i], M_unrot[i] = layer(h, M[i], M_unrot[i], pos) h = self.final_norm(h) logits = h @ self.embed_tokens.weight.T if self.config.tie_word_embeddings else self.lm_head(h) return logits, M, M_unrot