| """ |
| 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 |
|
|
