| """Parameter-audited Nero-XS-2 candidates for PyTorch/XLA.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import json |
| import math |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass(frozen=True) |
| class NeroConfig: |
| vocab_size: int = 2048 |
| width: int = 128 |
| heads: int = 4 |
| kv_heads: int = 2 |
| stored_blocks: int = 10 |
| ffn_width: int = 531 |
| recurrent_start: int = 1 |
| recurrent_blocks: int = 4 |
| recurrent_passes: int = 2 |
| engram_entries: int = 768 |
| use_engram: bool = True |
| use_qk_norm: bool = True |
| use_loop_conditioning: bool = True |
| max_position_embeddings: int = 2048 |
| rope_theta: float = 20000.0 |
|
|
| @property |
| def head_dim(self) -> int: |
| return self.width // self.heads |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, width: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(width)) |
| self.eps = eps |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight |
|
|
|
|
| class CenteredUnitNorm(nn.Module): |
| def __init__(self, width: int, eps: float = 1e-5): |
| super().__init__() |
| self.scale = nn.Parameter(torch.ones(width)) |
| self.shift = nn.Parameter(torch.zeros(width)) |
| self.eps = eps |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| centered = x - x.mean(-1, keepdim=True) |
| return centered * torch.rsqrt(centered.square().mean(-1, keepdim=True) + self.eps) * self.scale + self.shift |
|
|
|
|
| def deterministic_coordinates(length: int, width: int, base: float, device, dtype): |
| half = (width + 1) // 2 |
| positions = torch.arange(length, device=device, dtype=torch.float32)[:, None] |
| frequencies = torch.exp(torch.arange(half, device=device, dtype=torch.float32) * (-math.log(base) / max(half - 1, 1))) |
| result = torch.cat((torch.sin(positions * frequencies), torch.cos(positions * frequencies)), dim=-1)[:, :width] |
| return result.to(dtype) |
|
|
|
|
| def apply_rope(x: torch.Tensor, theta: float) -> torch.Tensor: |
| _, _, length, dim = x.shape |
| inv = theta ** (-torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim) |
| angles = torch.arange(length, device=x.device, dtype=torch.float32)[:, None] * inv[None, :] |
| cos = angles.cos().to(x.dtype)[None, None, :, :] |
| sin = angles.sin().to(x.dtype)[None, None, :, :] |
| even, odd = x[..., ::2], x[..., 1::2] |
| return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2) |
|
|
|
|
| class EngramLite(nn.Module): |
| """Collision-tolerant bigram/trigram memory with a contextual read gate.""" |
|
|
| def __init__(self, cfg: NeroConfig): |
| super().__init__() |
| self.entries = cfg.engram_entries |
| self.tables = nn.ModuleList([nn.Embedding(cfg.engram_entries, cfg.width) for _ in range(2)]) |
| self.gate = nn.Linear(cfg.width, 2, bias=True) |
| self.scale = nn.Parameter(torch.tensor(0.1)) |
|
|
| def _hash(self, ids: torch.Tensor, order: int, prime: int) -> torch.Tensor: |
| padded = F.pad(ids, (order - 1, 0), value=0) |
| value = torch.zeros_like(ids) |
| for offset in range(order): |
| value = (value * prime + padded[:, offset : offset + ids.shape[1]]) % self.entries |
| return value |
|
|
| def forward(self, ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor: |
| bigram = self.tables[0](self._hash(ids, 2, 10007)) |
| trigram = self.tables[1](self._hash(ids, 3, 10009)) |
| weights = torch.sigmoid(self.gate(hidden)) |
| memory = weights[..., :1] * bigram + weights[..., 1:] * trigram |
| return hidden + self.scale.tanh() * memory |
|
|
|
|
| class XSAAttention(nn.Module): |
| def __init__(self, cfg: NeroConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.q = nn.Linear(cfg.width, cfg.heads * cfg.head_dim, bias=False) |
| self.k = nn.Linear(cfg.width, cfg.kv_heads * cfg.head_dim, bias=False) |
| self.v = nn.Linear(cfg.width, cfg.kv_heads * cfg.head_dim, bias=False) |
| self.o = nn.Linear(cfg.heads * cfg.head_dim, cfg.width, bias=False) |
| self.q_norm = RMSNorm(cfg.head_dim) if cfg.use_qk_norm else nn.Identity() |
| self.k_norm = RMSNorm(cfg.head_dim) if cfg.use_qk_norm else nn.Identity() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| batch, length, _ = x.shape |
| q = self.q(x).view(batch, length, self.cfg.heads, self.cfg.head_dim).transpose(1, 2) |
| k = self.k(x).view(batch, length, self.cfg.kv_heads, self.cfg.head_dim).transpose(1, 2) |
| v = self.v(x).view(batch, length, self.cfg.kv_heads, self.cfg.head_dim).transpose(1, 2) |
| q = apply_rope(self.q_norm(q), self.cfg.rope_theta) |
| k = apply_rope(self.k_norm(k), self.cfg.rope_theta) |
| groups = self.cfg.heads // self.cfg.kv_heads |
| if groups > 1: |
| k = k.repeat_interleave(groups, dim=1) |
| v = v.repeat_interleave(groups, dim=1) |
| attended = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| unit_v = F.normalize(v, p=2, dim=-1, eps=1e-6) |
| attended = attended - (attended * unit_v).sum(-1, keepdim=True) * unit_v |
| return self.o(attended.transpose(1, 2).contiguous().view(batch, length, -1)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg: NeroConfig): |
| super().__init__() |
| self.attn_norm = RMSNorm(cfg.width) |
| self.attn = XSAAttention(cfg) |
| self.ffn_norm = RMSNorm(cfg.width) |
| self.gate = nn.Linear(cfg.width, cfg.ffn_width, bias=False) |
| self.up = nn.Linear(cfg.width, cfg.ffn_width, bias=False) |
| self.down = nn.Linear(cfg.ffn_width, cfg.width, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.attn_norm(x)) |
| normed = self.ffn_norm(x) |
| return x + self.down(F.silu(self.gate(normed)) * self.up(normed)) |
|
|
|
|
| class ReleasedXSAAttention(nn.Module): |
| def __init__(self, width: int = 128, heads: int = 4): |
| super().__init__() |
| self.width, self.heads, self.head_dim = width, heads, width // heads |
| self.q = nn.Linear(width, width, bias=False) |
| self.k = nn.Linear(width, width, bias=False) |
| self.v = nn.Linear(width, width, bias=False) |
| self.o = nn.Linear(width, width, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| batch, length, _ = x.shape |
| split = lambda value: value.view(batch, length, self.heads, self.head_dim).transpose(1, 2) |
| q, k, v = split(self.q(x)), split(self.k(x)), split(self.v(x)) |
| attended = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| coefficient = (attended * v).sum(-1, keepdim=True) / v.square().sum(-1, keepdim=True).clamp_min(1e-6) |
| attended = attended - coefficient * v |
| return self.o(attended.transpose(1, 2).contiguous().view(batch, length, self.width)) |
|
|
|
|
| class ReleasedBlock(nn.Module): |
| def __init__(self, width: int = 128, ffn_width: int = 540): |
| super().__init__() |
| self.attn_norm = CenteredUnitNorm(width) |
| self.attn = ReleasedXSAAttention(width) |
| self.ffn_norm = CenteredUnitNorm(width) |
| self.expand = nn.Linear(width, 2 * ffn_width, bias=False) |
| self.contract = nn.Linear(ffn_width, width, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.attn_norm(x)) |
| content, gate = self.expand(self.ffn_norm(x)).chunk(2, dim=-1) |
| return x + self.contract(F.silu(content) * torch.sigmoid(gate)) |
|
|
|
|
| class ReleasedNeroXSControl(nn.Module): |
| """Faithful PyTorch control for the released 2,996,480-parameter graph.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.config = NeroConfig(kv_heads=4, ffn_width=540, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False) |
| self.embedding = nn.Embedding(2048, 128) |
| self.blocks = nn.ModuleList([ReleasedBlock() for _ in range(10)]) |
| self.norm = CenteredUnitNorm(128) |
|
|
| def forward(self, ids: torch.Tensor) -> torch.Tensor: |
| x = self.embedding(ids) + deterministic_coordinates(ids.shape[1], 128, 20000.0, ids.device, self.embedding.weight.dtype)[None] |
| x = self.blocks[0](x) |
| for _ in range(2): |
| for block in self.blocks[1:5]: |
| x = block(x) |
| for block in self.blocks[5:]: |
| x = block(x) |
| return F.linear(self.norm(x), self.embedding.weight) |
|
|
|
|
| class NeroXSA2ForCausalLM(nn.Module): |
| """XSA + selective recurrence + loop conditioning + EngramLite.""" |
|
|
| def __init__(self, cfg: NeroConfig = NeroConfig()): |
| super().__init__() |
| self.config = cfg |
| self.embedding = nn.Embedding(cfg.vocab_size, cfg.width) |
| self.engram = EngramLite(cfg) if cfg.use_engram else None |
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.stored_blocks)]) |
| self.loop_embeddings = nn.Parameter(torch.zeros(cfg.recurrent_passes, cfg.width)) |
| self.loop_gates = nn.Parameter(torch.zeros(cfg.recurrent_passes, cfg.recurrent_blocks, cfg.width)) |
| self.norm = RMSNorm(cfg.width) |
| nn.init.normal_(self.loop_embeddings, std=0.01) |
|
|
| def forward(self, ids: torch.Tensor) -> torch.Tensor: |
| x = self.embedding(ids) |
| if self.engram is not None: |
| x = self.engram(ids, x) |
| start = self.config.recurrent_start |
| stop = start + self.config.recurrent_blocks |
| for block in self.blocks[:start]: |
| x = block(x) |
| for pass_index in range(self.config.recurrent_passes): |
| if self.config.use_loop_conditioning: |
| x = x + self.loop_embeddings[pass_index] |
| for local_index, block in enumerate(self.blocks[start:stop]): |
| if self.config.use_loop_conditioning: |
| proposal = block(x) |
| gate = torch.sigmoid(self.loop_gates[pass_index, local_index])[None, None, :] |
| x = x + gate * (proposal - x) |
| else: |
| x = block(x) |
| for block in self.blocks[stop:]: |
| x = block(x) |
| return F.linear(self.norm(x), self.embedding.weight) |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
|
|
| def variant_config(name: str) -> NeroConfig: |
| if name == "control": |
| return NeroConfig(kv_heads=4, ffn_width=540, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False) |
| if name == "gqa": |
| return NeroConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=False, use_loop_conditioning=False) |
| if name == "gqa_qknorm": |
| return NeroConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=True, use_loop_conditioning=False) |
| if name == "loop_conditioned": |
| return NeroConfig(kv_heads=2, ffn_width=582, engram_entries=0, use_engram=False, use_qk_norm=True, use_loop_conditioning=True) |
| if name == "full": |
| return NeroConfig() |
| raise ValueError(f"unknown architecture variant: {name}") |
|
|
|
|
| def build_variant(name: str) -> nn.Module: |
| if name == "released_control": |
| return ReleasedNeroXSControl() |
| return NeroXSA2ForCausalLM(variant_config(name)) |
|
|
|
|
| def architecture_audit() -> dict[str, int | float | bool]: |
| results = {} |
| for name in ("released_control", "control", "gqa", "gqa_qknorm", "loop_conditioned", "full"): |
| model = build_variant(name) |
| cfg = model.config |
| count = parameter_count(model) |
| if count >= 3_000_000: |
| raise AssertionError(f"{name} exceeds the parameter cap: {count:,}") |
| ids = torch.arange(64).remainder(cfg.vocab_size).view(1, -1) |
| with torch.no_grad(): |
| logits = model(ids) |
| changed = ids.clone() |
| changed[:, 32:] = (changed[:, 32:] + 17) % cfg.vocab_size |
| changed_logits = model(changed) |
| prefix_error = float((logits[:, :32] - changed_logits[:, :32]).abs().max()) |
| if logits.shape != (1, 64, cfg.vocab_size) or prefix_error > 1e-5: |
| raise AssertionError((name, logits.shape, prefix_error)) |
| results[name] = {"parameters": count, "causal_prefix_max_error": prefix_error} |
| return results |
|
|
|
|
| def load_model(model_dir: str | Path, device: str | torch.device = "cpu"): |
| """Load the released safetensors checkpoint and tokenizer.""" |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| model_dir = Path(model_dir) |
| raw = json.loads((model_dir / "config.json").read_text()) |
| fields = NeroConfig.__dataclass_fields__ |
| cfg = NeroConfig(**{key: value for key, value in raw.items() if key in fields}) |
| model = NeroXSA2ForCausalLM(cfg) |
| model.load_state_dict(load_file(model_dir / "model.safetensors"), strict=True) |
| model.to(device).eval() |
| tokenizer = AutoTokenizer.from_pretrained(model_dir) |
| return model, tokenizer |
|
|
|
|
| @torch.inference_mode() |
| def generate( |
| model: NeroXSA2ForCausalLM, |
| tokenizer, |
| prompt: str, |
| max_new_tokens: int = 64, |
| temperature: float = 0.8, |
| top_p: float = 0.95, |
| repetition_penalty: float = 1.1, |
| seed: int = 7, |
| ) -> str: |
| """Simple deterministic-seed nucleus sampler using full-prefix recomputation.""" |
| device = model.embedding.weight.device |
| ids = tokenizer.encode(prompt, add_special_tokens=False) |
| generator = torch.Generator(device=device).manual_seed(seed) |
| for _ in range(max_new_tokens): |
| context = torch.tensor([ids[-model.config.max_position_embeddings :]], device=device) |
| logits = model(context)[0, -1].float() |
| if repetition_penalty != 1.0: |
| seen = torch.tensor(sorted(set(ids)), device=device) |
| logits[seen] = torch.where( |
| logits[seen] < 0, |
| logits[seen] * repetition_penalty, |
| logits[seen] / repetition_penalty, |
| ) |
| if temperature <= 0: |
| next_id = int(logits.argmax()) |
| else: |
| sorted_logits, sorted_ids = (logits / temperature).sort(descending=True) |
| probabilities = sorted_logits.softmax(-1) |
| keep = probabilities.cumsum(-1) <= top_p |
| keep[0] = True |
| filtered = probabilities * keep |
| choice = torch.multinomial(filtered / filtered.sum(), 1, generator=generator) |
| next_id = int(sorted_ids[choice]) |
| ids.append(next_id) |
| if next_id == tokenizer.eos_token_id: |
| break |
| return tokenizer.decode(ids, skip_special_tokens=True) |
|
|
|
|
| if __name__ == "__main__": |
| import json |
| print("NERO_XS2_AUDIT=" + json.dumps(architecture_audit(), sort_keys=True)) |
|
|
