mid-run checkpoint step 8000 (0.26B tokens)

README.md

@@ -0,0 +1,100 @@
+---
+license: gpl-3.0
+library_name: model
+tags:
+- state-space-model
+- recurrent
+- language-model
+- from-scratch
+- plastic
+- mamba-like
+- triton
+datasets:
+- roneneldan/TinyStories
+- HuggingFaceFW/fineweb-edu
+metrics:
+- perplexity
+- accuracy
+---
+
+# pssa-100m
+
+A 109M-parameter language model with no attention, trained from scratch in a live notebook on one RTX PRO 6000 Blackwell. The [plastic state-space architecture](https://github.com/Sparticle62ops/pssa) (PSSA: selective diagonal state-space scan, Poincaré-ball episodic memory read, gated adapter, SiLU MLP) reimplemented in PyTorch, so the architecture's first 100M-scale checkpoint exists as plain weights you can load and sample in a few lines.
+
+The Rust repo's headline result is a 1.5M-parameter PSSA beating a matched transformer by 0.45 nats on WikiText-103. This checkpoint is the architecture's first answer at 109M: 2.5B tokens in three runs on one GPU, loss curves and sample text below. What you get is a toy that writes real English in two registers. What you don't get is facts. At 109M params it has capacity for style, not knowledge, and the web text it emits is confidently invented.
+
+## Architecture
+
+18 stacked blocks of:
+
+```
+xn       = rmsnorm(x)
+delta    = softplus(W_delta xn)         per-channel step size
+h_t      = exp(delta_t·A) h_{t-1} + delta_t·B_t·x_t     diagonal SSM scan, A < 0
+y_t      = <h_t, C_t>
+q        = poincare(W_qx xn + W_qh y)   query from token AND state
+w        = softmax(-d_H(q, k)/tau)      hyperbolic distance to 512 memory slots
+m        = w @ v                        bounded episodic read
+z        = y/√d_s + gate·W_proj m + rank16-silu-adapter(xn) + MLP
+h_i      = h_{i-1} + block_i(h)/√depth
+```
+
+Parameter counts: 101.50M core (d_model 768, untied 16k-token embedding and head) plus 7.37M across the 18 memory banks. There is no attention anywhere in the model. Cost per token is constant in context length; the state is fixed-size.
+
+Two deviations from the Rust reference, both forced by scale: the memory bank is gradient-trained rather than surprise-gated (the reference's 3.5-nat novelty threshold stops firing within ~300 steps once the model is past the basics, so the write path would sit inert for the rest of a 2.5B-token run), and backprop is truncated at 256 tokens. Training used a fused Triton scan (forward + backward in one kernel each) that matches a plain PyTorch loop to ~1e-7 on every gradient.
+
+## Training
+
+Three runs, one GPU, always resuming the same checkpoint:
+
+| run | data | tokens | speed |
+|---|---|---|---|
+| 1 | TinyStories only | 400M | 71k tok/s |
+| 2 | 50/50 TinyStories / FineWeb-Edu | +1.1B | 128k tok/s |
+| 3 | fine-tune, 85% FineWeb-Edu | +1.0B | 128k tok/s |
+
+The speed jump between runs 1 and 2 is `torch.compile` on the block: Inductor fuses ~40 small elementwise ops (the Poincaré distances, memory softmax, gates) into a handful of kernels, 184→130 ms/step, gradient cosine 0.9998 vs eager. Batch went 32→64 and chunk 128→256 on top of that. 99% GPU utilization held for the whole of runs 2 and 3.
+
+## Results
+
+Held-out, 1M unseen tokens per corpus:
+
+| | cross-entropy | perplexity | next-token acc |
+|---|---|---|---|
+| TinyStories | 1.329 | 3.78 | 64.5% |
+| FineWeb-Edu web | 2.935 | 18.8 | 43.4% |
+
+The interesting datapoint is run 2: moving from story-only to a 50/50 mixture *improved* story loss (1.302 → 1.244) while learning web text at all. Story loss only gave back ground when the mixture went to 85% web.
+
+Samples, temp 0.7 top-k 40:
+
+> Once upon a time there was a tiny dragon. The dragon was very fast and loved to run. He had a big, soft voice and a long tail. He lived in a deep, dark cave. One day, the dragon heard a voice. It was a little bird. The bird was lost and scared...
+
+> In the summer of 1969, the U.S. Department of Agriculture announced that it would implement "universally mandated standards for the development of new products and products." In the first phase of its series of "Design, Material and Technologies", the Standards for Energy and Smart Buildings (SEP) was developed and funded...
+
+Web loss was still falling when the budget ran out. This model is underfit, not overfit; more tokens would keep helping.
+
+## Use
+
+```python
+# pip install torch safetensors tokenizers
+from pssa_model import from_pretrained, generate
+
+model, tok = from_pretrained(".")        # dir containing model.safetensors + tokenizer.json
+print(generate(model, tok, "Once upon a time", max_new=100))
+```
+
+`pssa_model.py` is self-contained: `PSSAConfig`, `PSSAModel`, loader, sampler. On CPU it runs a plain PyTorch scan; on CUDA it uses the fused Triton kernel from `pssa_kernels.py` when that file is importable. No transformers, no configuration — plain weights and ~150 lines.
+
+Weights are bf16 safetensors, 217 MB. The 16,384-token byte-level BPE tokenizer trained on TinyStories ships alongside. It handles general web text without unmangling, at some cost in efficiency on rare Unicode.
+
+## Files
+
+- `model.safetensors`: weights, bf16
+- `pssa_model.py`: model + loader + sampler
+- `pssa_kernels.py`: fused Triton scan
+- `tokenizer.json`: 16k byte-level BPE
+
+## Credit
+
+Architecture, reference implementation and the small-scale comparison: [Sparticle62ops/pssa](https://github.com/Sparticle62ops/pssa) (GPL-3.0). This checkpoint was trained in a marimo notebook on a single RTX PRO 6000 Blackwell; the training notebook is in the repo's `notebooks/` directory.

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:751607c8f8f2dc9d30a431b902ca20d5219407d429046811af493d451ada567b
+size 217776144

pssa_kernels.py

@@ -0,0 +1,151 @@
+"""Fused Triton scan for the PSSA selective diagonal state-space recurrence.
+
+    h_t = exp(delta_t * A) * h_{t-1} + delta_t * b_t * x_t
+    y_t = <h_t, c_t>
+
+`A` is (D, DS) negative, `delta` is (B, T, D) positive, `b`/`c` are (B, T, DS).
+The whole (forward + backward) runs as one kernel launch per chunk, with the
+per-token state `h_t` materialised so the backward pass is a single reverse
+sweep.  Carries `h0` across chunks and returns `hT`.
+"""
+import torch
+import triton
+import triton.language as tl
+
+__all__ = ["ScanFn", "scan_reference", "selftest"]
+
+
+@triton.jit
+def scan_fwd(x_ptr, d_ptr, bp_ptr, cp_ptr, A_ptr, h0_ptr, y_ptr, hT_ptr, H_ptr,
+             T: tl.constexpr, D: tl.constexpr, DS: tl.constexpr, BD: tl.constexpr):
+    b = tl.program_id(0)
+    pd = tl.program_id(1)
+    od = pd * BD + tl.arange(0, BD)
+    os_ = tl.arange(0, DS)
+    A = tl.load(A_ptr + od[:, None] * DS + os_[None, :])
+    h = tl.load(h0_ptr + b * D * DS + od[:, None] * DS + os_[None, :])
+    for t in range(T):
+        x = tl.load(x_ptr + b * T * D + t * D + od).to(tl.float32)
+        dl = tl.load(d_ptr + b * T * D + t * D + od).to(tl.float32)
+        bp = tl.load(bp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
+        cp = tl.load(cp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
+        h = tl.exp(dl[:, None] * A) * h + dl[:, None] * bp[None, :] * x[:, None]
+        tl.store(y_ptr + b * T * D + t * D + od, tl.sum(h * cp[None, :], axis=1))
+        tl.store(H_ptr + b * T * D * DS + t * D * DS + od[:, None] * DS + os_[None, :], h)
+    tl.store(hT_ptr + b * D * DS + od[:, None] * DS + os_[None, :], h)
+
+
+@triton.jit
+def scan_bwd(dy_ptr, x_ptr, d_ptr, bp_ptr, cp_ptr, A_ptr, H_ptr, h0_ptr,
+             gh0_ptr, gx_ptr, gd_ptr, gbp_ptr, gcp_ptr, gA_ptr,
+             T: tl.constexpr, D: tl.constexpr, DS: tl.constexpr, BD: tl.constexpr):
+    b = tl.program_id(0)
+    pd = tl.program_id(1)
+    od = pd * BD + tl.arange(0, BD)
+    os_ = tl.arange(0, DS)
+    A = tl.load(A_ptr + od[:, None] * DS + os_[None, :])
+    g = tl.zeros((BD, DS), tl.float32)
+    gA = tl.zeros((BD, DS), tl.float32)
+    for t in range(T - 1, -1, -1):
+        dy = tl.load(dy_ptr + b * T * D + t * D + od).to(tl.float32)
+        x = tl.load(x_ptr + b * T * D + t * D + od).to(tl.float32)
+        dl = tl.load(d_ptr + b * T * D + t * D + od).to(tl.float32)
+        bp = tl.load(bp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
+        cp = tl.load(cp_ptr + b * T * DS + t * DS + os_).to(tl.float32)
+        h = tl.load(H_ptr + b * T * D * DS + t * D * DS + od[:, None] * DS + os_[None, :])
+        if t < T - 1:
+            dln = tl.load(d_ptr + b * T * D + (t + 1) * D + od).to(tl.float32)
+            g = dy[:, None] * cp[None, :] + tl.exp(dln[:, None] * A) * g
+        else:
+            g = dy[:, None] * cp[None, :]
+        if t > 0:
+            hp = tl.load(H_ptr + b * T * D * DS + (t - 1) * D * DS + od[:, None] * DS + os_[None, :])
+        else:
+            hp = tl.load(h0_ptr + b * D * DS + od[:, None] * DS + os_[None, :])
+        a = tl.exp(dl[:, None] * A)
+        ga = g * hp
+        gd = tl.sum(ga * a * A, axis=1) + tl.sum(g * bp[None, :] * x[:, None], axis=1)
+        gbp = tl.sum(g * dl[:, None] * x[:, None], axis=0)
+        gcp = tl.sum(dy[:, None] * h, axis=0)
+        gx = tl.sum(g * dl[:, None] * bp[None, :], axis=1)
+        gA += ga * a * dl[:, None]
+        if t == 0:
+            tl.store(gh0_ptr + b * D * DS + od[:, None] * DS + os_[None, :], a * g)
+        tl.store(gx_ptr + b * T * D + t * D + od, gx)
+        tl.store(gd_ptr + b * T * D + t * D + od, gd)
+        tl.atomic_add(gbp_ptr + b * T * DS + t * DS + os_, gbp)
+        tl.atomic_add(gcp_ptr + b * T * DS + t * DS + os_, gcp)
+    tl.atomic_add(gA_ptr + od[:, None] * DS + os_[None, :], gA)
+
+
+class ScanFn(torch.autograd.Function):
+    """Differentiable fused scan. Returns (y, hT); hT is not differentiated."""
+
+    @staticmethod
+    def forward(ctx, x, delta, bp, cp, A, h0, bd=16):
+        B, T, D = x.shape
+        DS = A.shape[1]
+        y = torch.empty(B, T, D, device=x.device, dtype=torch.float32)
+        hT = torch.empty(B, D, DS, device=x.device, dtype=torch.float32)
+        H = torch.empty(B, T, D, DS, device=x.device, dtype=torch.float32)
+        scan_fwd[(B, triton.cdiv(D, bd))](
+            x, delta, bp, cp, A, h0, y, hT, H, T=T, D=D, DS=DS, BD=bd)
+        ctx.save_for_backward(x, delta, bp, cp, A, H, h0)
+        ctx.bd = bd
+        ctx.mark_non_differentiable(hT)
+        return y, hT
+
+    @staticmethod
+    def backward(ctx, dy, dhT):
+        x, delta, bp, cp, A, H, h0 = ctx.saved_tensors
+        B, T, D = x.shape
+        DS = A.shape[1]
+        bd = ctx.bd
+        gx = torch.empty(B, T, D, device=x.device, dtype=torch.float32)
+        gd = torch.empty_like(delta, dtype=torch.float32)
+        gbp = torch.zeros_like(bp, dtype=torch.float32)
+        gcp = torch.zeros_like(cp, dtype=torch.float32)
+        gA = torch.zeros_like(A, dtype=torch.float32)
+        gh0 = torch.zeros(B, D, DS, device=x.device, dtype=torch.float32)
+        scan_bwd[(B, triton.cdiv(D, bd))](
+            dy.contiguous(), x, delta, bp, cp, A, H, h0, gh0,
+            gx, gd, gbp, gcp, gA, T=T, D=D, DS=DS, BD=bd)
+        return gx.to(x.dtype), gd, gbp, gcp, gA, gh0.to(h0.dtype), None
+
+
+def scan_reference(x, delta, bp, cp, A, h0):
+    """Plain PyTorch loop, for tests. Returns (y, H, hT)."""
+    B, T, D = x.shape
+    h = h0
+    ys, hs = [], []
+    for t in range(T):
+        h = torch.exp(delta[:, t, :, None] * A) * h + delta[:, t, :, None] * bp[:, t, None, :] * x[:, t, :, None]
+        ys.append((h * cp[:, t, None, :]).sum(-1))
+        hs.append(h)
+    return torch.stack(ys, 1), torch.stack(hs, 1), h
+
+
+def selftest(ds=16, d=64, b=2, t=7, verbose=True):
+    torch.manual_seed(1)
+    dev = "cuda"
+    x = torch.randn(b, t, d, device=dev, requires_grad=True)
+    delta = torch.nn.functional.softplus(torch.randn(b, t, d, device=dev)).requires_grad_(True)
+    bp = torch.randn(b, t, ds, device=dev, requires_grad=True)
+    cp = torch.randn(b, t, ds, device=dev, requires_grad=True)
+    A = (-torch.nn.functional.softplus(torch.randn(d, ds, device=dev))).requires_grad_(True)
+    h0 = (torch.randn(b, d, ds, device=dev) * 0.5).requires_grad_(True)
+    yr, _, _ = scan_reference(x, delta, bp, cp, A, h0)
+    y, hT = ScanFn.apply(x, delta, bp, cp, A, h0)
+    gy = torch.randn_like(yr)
+    gref = torch.autograd.grad(yr, (x, delta, bp, cp, A, h0), gy)
+    gtri = torch.autograd.grad(y, (x, delta, bp, cp, A, h0), gy)
+    out = {"fwd": ((y - yr).abs().max() / yr.abs().max()).item()}
+    for nm, a, bb in zip(["dx", "dd", "dbp", "dcp", "dA", "dh0"], gtri, gref):
+        out[nm] = ((a - bb).abs().max() / bb.abs().max()).item()
+    if verbose:
+        print("scan selftest rel errs: " + "  ".join(f"{k}={v:.1e}" for k, v in out.items()))
+    return out
+
+
+if __name__ == "__main__":
+    selftest()

pssa_model.py

@@ -0,0 +1,179 @@
+"""Standalone PSSA model for the Hub: numpy-only inference, bf16 safetensors."""
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+try:
+    from pssa_kernels import ScanFn
+except Exception:  # no compiled kernel (CPU-only install): fall back below
+    ScanFn = None
+
+
+def _scan(x, delta, bp, cp, A, h0):
+    """h_t = exp(delta_t·A) h_{t-1} + delta_t·b_t·x_t, y_t = <h_t, c_t>."""
+    if ScanFn is not None and x.is_cuda:
+        return ScanFn.apply(x, delta, bp, cp, A, h0)
+    B, T, D = x.shape
+    h = h0
+    ys = []
+    for t in range(T):
+        h = torch.exp(delta[:, t, :, None] * A) * h + delta[:, t, :, None] * bp[:, t, None, :] * x[:, t, :, None]
+        ys.append((h * cp[:, t, None, :]).sum(-1))
+    return torch.stack(ys, 1), h
+
+
+@dataclass
+class PSSAConfig:
+    vocab_size: int = 16384
+    d_model: int = 768
+    depth: int = 18
+    d_state: int = 16
+    d_key: int = 32
+    mem_slots: int = 512
+    adapter_rank: int = 16
+    mlp_hidden: int = 1536
+    chunk: int = 256
+    seq_len: int = 512
+    tau_mem: float = 0.1
+    value_cap: float = 8.0
+
+
+def _xavier(m):
+    if isinstance(m, nn.Linear):
+        bound = math.sqrt(6 / (m.in_features + m.out_features))
+        nn.init.uniform_(m.weight, -bound, bound)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+
+
+def _poincare(x):
+    """q -> q / (1 + |q|): maps R^n into the open unit ball."""
+    return x / (1.0 + x.norm(dim=-1, keepdim=True))
+
+
+class PSSABlock(nn.Module):
+    """SSM scan -> hyperbolic memory read -> gate -> plastic adapter -> MLP."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        d, ds, dk, r = cfg.d_model, cfg.d_state, cfg.d_key, cfg.adapter_rank
+        self.cfg = cfg
+        self.norm = nn.Parameter(torch.ones(d))
+        self.norm_b = nn.Parameter(torch.zeros(d))
+        self.w_delta = nn.Linear(d, d, bias=False)
+        self.w_b = nn.Linear(d, ds, bias=False)
+        self.w_c = nn.Linear(d, ds, bias=False)
+        ratio = torch.arange(ds).float() / (ds - 1)
+        tau = torch.exp(math.log(1.5) + ratio * (math.log(200.0) - math.log(1.5)))
+        self.a_raw = nn.Parameter(torch.log(torch.expm1(1.0 / tau)).unsqueeze(0).repeat(d, 1))
+        self.w_qx = nn.Linear(d, dk, bias=False)
+        self.w_qh = nn.Linear(d, dk, bias=False)
+        self.w_gate = nn.Linear(d, d, bias=False)
+        self.w_proj = nn.Linear(d, d, bias=False)
+        self.ad_down = nn.Linear(d, r, bias=False)
+        self.ad_up = nn.Linear(r, d, bias=False)
+        self.mlp1 = nn.Linear(d, cfg.mlp_hidden, bias=False)
+        self.mlp2 = nn.Linear(cfg.mlp_hidden, d, bias=False)
+        self.mk_raw = nn.Parameter(torch.randn(cfg.mem_slots, dk) * 0.05)
+        self.mv = nn.Parameter(torch.randn(cfg.mem_slots, d) * 0.02)
+        for m in (self.w_delta, self.w_b, self.w_c, self.w_qx, self.w_qh,
+                  self.w_gate, self.ad_down, self.mlp1):
+            _xavier(m)
+        nn.init.zeros_(self.w_proj.weight)
+        nn.init.zeros_(self.ad_up.weight)
+        nn.init.zeros_(self.mlp2.weight)
+
+    def forward(self, x, h0):
+        cfg = self.cfg
+        xn = self.norm * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5)) + self.norm_b
+        delta = F.softplus(self.w_delta(xn).float())
+        bp = self.w_b(xn).float()
+        cp = self.w_c(xn).float()
+        A = -F.softplus(self.a_raw)
+        y, h = _scan(xn, delta, bp, cp, A, h0)
+
+        q = _poincare((self.w_qx(xn) + self.w_qh(y)).float())
+        mk = _poincare(self.mk_raw)
+        qsq = (q * q).sum(-1, keepdim=True)
+        mksq = (mk * mk).sum(-1)
+        sq = (qsq + mksq.view(1, 1, -1) - 2 * (q @ mk.T)).clamp_min(0)
+        den = ((1 - qsq) * (1 - mksq).view(1, 1, -1)).clamp_min(1e-12)
+        dist = 2 * torch.asinh(torch.sqrt(sq / den))
+        w = torch.softmax(-dist / cfg.tau_mem, -1)
+        mv = cfg.value_cap * self.mv / (cfg.value_cap + self.mv.norm(dim=-1, keepdim=True))
+        m = w @ mv
+
+        g = torch.sigmoid(self.w_gate(xn))
+        z = y / math.sqrt(cfg.d_state) + g * self.w_proj(m.to(x.dtype))
+        z = z + self.ad_up(F.silu(self.ad_down(xn)))
+        z = z + self.mlp2(F.silu(self.mlp1(z)))
+        return z, h
+
+
+class PSSAModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.blocks = nn.ModuleList([PSSABlock(cfg) for _ in range(cfg.depth)])
+        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.scale = 1.0 / math.sqrt(cfg.depth)
+
+    def init_state(self, batch_size, device=None):
+        dev = device or self.embed.weight.device
+        return [torch.zeros(batch_size, self.cfg.d_model, self.cfg.d_state, device=dev)
+                for _ in range(self.cfg.depth)]
+
+    def forward_step(self, tokens, states):
+        x = self.embed(tokens)
+        for i, blk in enumerate(self.blocks):
+            z, h = blk(x, states[i])
+            x = z if i == 0 else x + z * self.scale
+            states[i] = h
+        return self.head(x) / math.sqrt(self.cfg.d_model), states
+
+    def forward(self, tokens):
+        states = self.init_state(tokens.shape[0], tokens.device)
+        outs = []
+        for s in range(0, tokens.shape[1], self.cfg.chunk):
+            logits, states = self.forward_step(tokens[:, s:s + self.cfg.chunk], states)
+            outs.append(logits)
+        return torch.cat(outs, 1)
+
+
+def from_pretrained(model_dir: str, device: str = "cpu"):
+    """Load the Hub export: (model, tokenizer). Works on CPU with no Triton install."""
+    from safetensors.torch import load_file
+    from tokenizers import Tokenizer
+
+    cfg = PSSAConfig()
+    sd = load_file(f"{model_dir}/model.safetensors")
+    sd = {k: v.float() for k, v in sd.items()}
+    m = PSSAModel(cfg)
+    m.load_state_dict(sd)
+    m.eval().to(device)
+    return m, Tokenizer.from_file(f"{model_dir}/tokenizer.json")
+
+
+@torch.no_grad()
+def generate(model, tokenizer, prompt, max_new=100, temp=0.7, top_k=40, eot=0, device="cpu"):
+    ids = tokenizer.encode(prompt).ids
+    states = model.init_state(1, device)
+    logits, states = model.forward_step(torch.tensor([ids], device=device), states)
+    out = []
+    for _ in range(max_new):
+        lg = logits[:, -1, :].float() / temp
+        if top_k:
+            kth = torch.topk(lg, top_k).values[:, [-1]]
+            lg = lg.masked_fill(lg < kth, -float("inf"))
+        nxt = torch.multinomial(torch.softmax(lg, -1), 1)
+        if int(nxt) == eot:
+            break
+        out.append(int(nxt))
+        logits, states = model.forward_step(nxt, states)
+    text = tokenizer.decode([i for i in out if i != eot])
+    return prompt + text