| """Architecture dimensions of the fixed inference export.""" |
| from dataclasses import dataclass |
|
|
|
|
| @dataclass(frozen=True) |
| class ModelConfig: |
| vocab_size: int = 2048 |
| width: int = 384 |
| heads: int = 6 |
| mlp_width: int = 2096 |
| loops: int = 6 |
| rank: int = 128 |
| controller_width: int = 16 |
| max_seq_len: int = 8192 |
| norm_eps: float = 1e-06 |
| routing_eps: float = 1.1920928955078125e-07 |
|
|
| def validate(self): |
| if self != ModelConfig(): |
| raise ValueError("This runtime supports the released Purrence architecture only") |
|
|
|
|
| """Six-loop Attention Residuals inference with per-head XSA and gates.""" |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
|
|
| def rms(x, eps): |
| work = x if x.dtype == torch.float64 else x.float() |
| return (work * torch.rsqrt(work.square().mean(-1, keepdim=True) + eps)).to(x.dtype) |
|
|
|
|
| def depth_read(sources, query, eps=2 ** -23): |
| values = torch.stack(sources) |
| work = values if values.dtype == torch.float64 else values.float() |
| inverse_rms = torch.rsqrt(work.square().mean(-1) + eps) |
| scores = (work[..., -query.numel():] * query.to(work.dtype)).sum(-1) * inverse_rms |
| weights = scores.softmax(dim=0) |
| return (weights[..., None] * work).sum(0).to(values.dtype) |
|
|
|
|
| def document_layout(ids, boundary_id): |
| flat = ids.reshape(-1) |
| starts = flat.eq(boundary_id) |
| starts[0] = True |
| if ids.ndim == 2: |
| starts[::ids.shape[1]] = True |
| offsets = starts.nonzero().flatten() |
| cu = torch.cat((offsets, offsets.new_tensor([flat.numel()]))).to(torch.int32) |
| positions = torch.arange(flat.numel(), device=ids.device) - torch.repeat_interleave(offsets, cu.diff().long()) |
| return cu, positions |
|
|
|
|
| def half_rope(x, positions, frequencies, factors=None): |
| rotary_dim = x.shape[-1] // 2 |
| if factors is None: |
| angle = positions.float()[:, None] * frequencies[None] |
| cosine, sine = angle.cos().to(x.dtype)[:, None], angle.sin().to(x.dtype)[:, None] |
| else: |
| cosine, sine = (factor.to(x.dtype) for factor in factors) |
| even, odd = x[..., :rotary_dim:2], x[..., 1:rotary_dim:2] |
| rotated = torch.stack((even * cosine + odd * sine, odd * cosine - even * sine), dim=-1) |
| return torch.cat((rotated.flatten(-2), x[..., rotary_dim:]), dim=-1) |
|
|
|
|
| def offset_keys(keys, positions): |
| rotary_dim = keys.shape[-1] // 2 |
| previous = torch.cat((keys[:1, :, rotary_dim:], keys[:-1, :, rotary_dim:])) |
| stationary = torch.where(positions[:, None, None] > 0, previous, keys[..., rotary_dim:]) |
| return torch.cat((keys[..., :rotary_dim], stationary), dim=-1) |
|
|
|
|
| def attention(q, k, v, cu): |
| result = [] |
| boundaries = cu.tolist() |
| for start, end in zip(boundaries, boundaries[1:]): |
| query, key, value = (item[start:end].transpose(0, 1)[None] for item in (q, k, v)) |
| out = F.scaled_dot_product_attention(query, key, value, is_causal=True) |
| result.append(out[0].transpose(0, 1)) |
| return torch.cat(result) |
|
|
|
|
| class SharedBlock(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.q = nn.Linear(config.width, config.width, bias=False) |
| self.k = nn.Linear(config.width, config.width, bias=False) |
| self.v = nn.Linear(config.width, config.width, bias=False) |
| self.o = nn.Linear(config.width, config.width, bias=False) |
| self.up = nn.Linear(config.width, config.mlp_width, bias=False) |
| self.down = nn.Linear(config.mlp_width, config.width, bias=False) |
| self.xsa = nn.Linear(config.controller_width, config.heads, bias=False) |
| self.gate = nn.Linear(config.controller_width, config.heads, bias=False) |
| head_dim = config.width // config.heads |
| self.register_buffer("frequencies", (1 / 1024) ** torch.linspace(0, 1, head_dim // 4), persistent=False) |
|
|
| def rotary_factors(self, positions, dtype): |
| angle = positions.float()[:, None] * self.frequencies[None] |
| return angle.cos().to(dtype)[:, None], angle.sin().to(dtype)[:, None] |
|
|
| def attend(self, x, cu, positions, rope_factors): |
| config = self.config |
| q, k, v = [F.linear(x, getattr(self, name).weight.to(x.dtype)).view( |
| -1, config.heads, config.width // config.heads) for name in ('q', 'k', 'v')] |
| q = half_rope(rms(q, config.routing_eps), positions, self.frequencies, rope_factors) |
| k = half_rope(rms(k, config.routing_eps), positions, self.frequencies, rope_factors) |
| k = offset_keys(k, positions) |
| y = attention(q, k, v, cu) |
| local = x[..., :config.controller_width] |
| alpha, gate = self.xsa(local), self.gate(local) |
| dot = (y * v).sum(-1, keepdim=True) |
| denominator = v.square().sum(-1, keepdim=True).clamp_min(1e-8) |
| y = y - alpha.tanh().to(y.dtype)[..., None] * (dot / denominator) * v |
| y = y * gate.sigmoid().to(y.dtype)[..., None] |
| return F.linear(y.contiguous().flatten(-2), self.o.weight.to(y.dtype)) |
|
|
| def mlp(self, x): |
| return self.down(self.up(x).relu().square()) |
|
|
|
|
| class RecurrentLM(nn.Module): |
| def __init__(self, config=None): |
| super().__init__() |
| self.config = config or ModelConfig() |
| self.config.validate() |
| self.embedding = nn.Embedding(self.config.vocab_size, self.config.width) |
| self.block = SharedBlock(self.config) |
| self.queries = nn.Parameter(torch.zeros(2 * self.config.loops, self.config.rank)) |
|
|
| def read(self, sources, query): |
| x = sources[0] if query is None else depth_read(sources, query, self.config.routing_eps) |
| return rms(x, self.config.norm_eps) |
|
|
| def forward(self, ids, cu_seqlens=None, positions=None): |
| if ids.ndim not in (1, 2) or ids.numel() == 0: |
| raise ValueError("Expected nonempty one- or two-dimensional token IDs") |
| if cu_seqlens is None: |
| cu_seqlens, positions = document_layout(ids, self.config.vocab_size - 1) |
| if positions is None: |
| raise ValueError("Explicit packed boundaries require explicit document positions") |
| if int(cu_seqlens.diff().max()) > self.config.max_seq_len: |
| raise ValueError("Document segment exceeds max_seq_len") |
| embedding = self.embedding(ids.reshape(-1)) |
| rope_factors = self.block.rotary_factors(positions, embedding.dtype) |
| sources = [embedding] |
| for loop in range(self.config.loops): |
| x = self.read(sources, None if loop == 0 else self.queries[2 * loop - 1]) |
| a = self.block.attend(x, cu_seqlens, positions, rope_factors) |
| x = self.read(sources + [a], self.queries[2 * loop]) |
| m = self.block.mlp(x) |
| sources.append(a + m) |
| x = self.read(sources, self.queries[-1]) |
| return F.linear(x, self.embedding.weight).view(*ids.shape, self.config.vocab_size) |
|
|
| def parameter_count(self): |
| return sum(parameter.numel() for parameter in self.parameters()) |
|
|
|
|
| |
| """Small runtime for the reviewed 2048-ID tokenizer and structural boundaries.""" |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
|
|
| class Sub3MTokenizer: |
| vocab_size = 2048 |
| boundary_id = 2047 |
|
|
| def __init__(self, artifact_dir): |
| import tokenizers |
| directory = Path(artifact_dir) |
| metadata = json.loads((directory / "artifact.json").read_text(encoding="utf-8")) |
| if tokenizers.__version__ != metadata["tokenizers_version"]: |
| raise RuntimeError("Use tokenizers==" + metadata["tokenizers_version"]) |
| path = directory / "ordinary-tokenizer.json" |
| if hashlib.sha256(path.read_bytes()).hexdigest() != metadata["tokenizer_sha256"]: |
| raise ValueError("Tokenizer file does not match its artifact hash") |
| self.backend = tokenizers.Tokenizer.from_file(str(path)) |
| if self.backend.get_vocab_size() != 2047 or metadata["vocab_size"] != 2048: |
| raise ValueError("Expected 2047 ordinary IDs and one structural boundary") |
| if json.loads(path.read_text(encoding="utf-8")).get("added_tokens"): |
| raise ValueError("Ordinary encoder must not recognize structural token spellings") |
| self.metadata = metadata |
|
|
| def encode(self, text): |
| """Encode literal text; even '<|endoftext|>' remains ordinary content.""" |
| if not isinstance(text, str): |
| raise TypeError("encode expects a Unicode string") |
| ids = self.backend.encode(text, add_special_tokens=False).ids |
| if any(token_id >= self.boundary_id for token_id in ids): |
| raise ValueError("Ordinary encoding emitted a reserved ID") |
| return ids |
|
|
| def decode(self, ids, skip_boundaries=False): |
| """Decode generated IDs; optionally omit structural boundary markers. |
| |
| Arbitrary generated byte sequences can be invalid UTF-8; the ByteLevel |
| decoder replaces invalid sequences. Every encoded valid string roundtrips. |
| """ |
| chunks = [] |
| ordinary = [] |
| for token_id in ids: |
| if type(token_id) is not int or token_id not in range(self.vocab_size): |
| raise ValueError("Token ID must be an integer in [0, 2048)") |
| if token_id == self.boundary_id: |
| chunks.append(self.backend.decode(ordinary, skip_special_tokens=False)) |
| ordinary = [] |
| if not skip_boundaries: |
| chunks.append(self.metadata["boundary_special"]["spelling"]) |
| else: |
| ordinary.append(token_id) |
| chunks.append(self.backend.decode(ordinary, skip_special_tokens=False)) |
| return "".join(chunks) |
|
|
