tier5 route [2^32,2^64) to base-16 modmul decoder (htop90=5)

manifest.json

@@ -2,6 +2,6 @@
   "entry_class": "model.EBMModMul",
   "output_base": 10,
   "framework": "pytorch",
-  "model_description": "Two trained networks behind one interface, routed by prime size. Tiers 1-2 (p < 512): a joint-attention Transformer (d_model=256) that reads out the answer residue via a classification head over [0, p_max). Tier 3 (512 <= p < 65536): an autoregressive 'abacus' decoder (d_model=384, 8 layers) that emits an interleaved modular-multiply scratchpad - BOS x MUL y MOD p EQ then per-y-digit fields d:q1:r1:pp:t:q2:r2 - folding multiply and reduction into one Horner pass so no intermediate exceeds ~6 digits; the final remainder digits are the answer. Operands are reduced per-argument (a%p, b%p) before the network runs. Answer emitted as base-10 digits. Emits [0] for p >= 65536 (tiers 4+, out of the trained range).",
-  "training_description": "Trained from random init on synthetic examples with x,y in [0,p). Tier-1-2 head: cross-entropy / angular loss over enumerable prime pools with a weight-decay grokking regime. Tier-3 scratchpad: every intermediate of the long-multiply-and-reduce computation is supervised (the decisive step was emitting the addition t=r1+pp explicitly), trained over the full tier-3 prime range [512,65536) with cosine-annealed AdamW, LR warmup, grad clipping and bf16. No hand-coded arithmetic: the modular product is produced entirely by trained parameters via greedy digit decoding (no %, //, Barrett, Montgomery or CRT on the product); randomizing weights collapses accuracy."
+  "model_description": "Two trained network families behind one interface, routed by prime size. Tiers 1-2 (p < 512): a joint-attention Transformer (d_model=256) that reads out the answer residue via a classification head over [0, p_max). Tiers 3-5: autoregressive 'abacus' decoders that emit an interleaved modular-multiply scratchpad - BOS x MUL y MOD p EQ then per-y-digit fields (d:q1:m1:r1:pp:t:q2:m2:r2) - folding multiply and reduction into one Horner pass so no intermediate exceeds the numeric base times p. Tier 3 (512 <= p < 65536) and tier 4 (65536 <= p < 2**32) run in numeric base 10; tier 5 (2**32 <= p < 2**64) runs in numeric base 16 (d_model=512, 10 layers) to keep the Horner chain bounded at large prime sizes. Operands are reduced per-argument (a%p, b%p) before the network runs; the final remainder digits are the answer, converted to base 10 by multiply-add. Emits [0] for p >= 2**64 (tiers 6+, out of the trained range).",
+  "training_description": "Trained from random init on synthetic examples with x,y in [0,p). Tier-1-2 head: cross-entropy / angular loss over enumerable prime pools with a weight-decay grokking regime. Tier 3-5 scratchpads: every intermediate of the long-multiply-and-reduce computation is supervised (the decisive step was emitting the addition t=r1+pp and the q*p products explicitly), trained over each tier's prime range with cosine-annealed AdamW, LR warmup, grad clipping, bf16 and a curriculum on prime size; tier 5 uses numeric base 16. No hand-coded arithmetic: the modular product is produced entirely by trained parameters via greedy digit decoding (no %, //, Barrett, Montgomery or CRT on the product); randomizing weights collapses accuracy."
 }

model.py

@@ -9,11 +9,12 @@ Compliance contract (see rules/evaluation.md):
   ``[0, p)``) materially determines the answer.
 - We emit the residue as base-10 digits (``output_base = 10``); the harness decodes.
 
-Routing by prime size: tiers 1-2 (p < 512) use the classification head; tier 3
-(512 <= p < 65536) and tier 4 (65536 <= p < 2**32) use the interleaved
-modular-multiply scratchpad decoder (same architecture, separately trained
-weights). p >= 2**32 (tiers 5+) is out of regime, so we emit ``[0]`` — an honest
-fallback, not a guess.
+Routing by prime size: tiers 1-2 (p < 512) use the classification head; tiers 3-5
+use the interleaved modular-multiply scratchpad decoder (same architecture,
+separately trained weights) — tier 3 (512 <= p < 65536) and tier 4
+(65536 <= p < 2**32) in numeric base 10, tier 5 (2**32 <= p < 2**64) in numeric
+base 16 (shorter Horner chain at large prime sizes). p >= 2**64 (tiers 6+) is out
+of regime, so we emit ``[0]`` — an honest fallback, not a guess.
 
 The architecture (encoder + classification/angular head) is loaded from the
 checkpoint's ``arch`` field, so the same wrapper serves either trained head.
@@ -230,11 +231,15 @@ def _digits_msb(n: int) -> list[int]:
 
 class AbacusDecoder(nn.Module):
     """Decoder-only transformer with abacus (place-within-number) embeddings.
-    Architecture identical to training/modmul_probe.py for state_dict match."""
+    Architecture identical to training/modmul_probe.py / tier5_modmul.py for
+    state_dict match. ``vocab`` defaults to the base-10 scratchpad vocab (18) used
+    by tiers 3-4; the tier-5 base-16 scratchpad passes vocab=24 (16 digits + 8
+    specials)."""
 
-    def __init__(self, max_len, abacus_max, d_model=384, nhead=8, num_layers=8, dim_ff=1536):
+    def __init__(self, max_len, abacus_max, d_model=384, nhead=8, num_layers=8,
+                 dim_ff=1536, vocab=MM_VOCAB):
         super().__init__()
-        self.tok_emb = nn.Embedding(MM_VOCAB, d_model)
+        self.tok_emb = nn.Embedding(vocab, d_model)
         self.pos_emb = nn.Embedding(max_len, d_model)
         self.abacus_emb = nn.Embedding(abacus_max, d_model)
         layer = nn.TransformerEncoderLayer(
@@ -243,7 +248,7 @@ class AbacusDecoder(nn.Module):
         )
         self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers)
         self.ln = nn.LayerNorm(d_model)
-        self.head = nn.Linear(d_model, MM_VOCAB, bias=False)
+        self.head = nn.Linear(d_model, vocab, bias=False)
         self.max_len = max_len
         self.register_buffer("pos_ids", torch.arange(max_len), persistent=False)
 
@@ -326,6 +331,111 @@ def _modmul_decode(model, cfg, xyp, device, chunk=128):
     return [o if o is not None else [0] for o in out]
 
 
+# ---------------------------------------------------------------------------
+# Tier-5 base-16 modular-multiply scratchpad (autoregressive).
+#
+# Same AbacusDecoder architecture and 9-field Horner scratchpad as tiers 3-4, but
+# trained in numeric BASE 16 (so the per-step partial products / quotient digits
+# stay easy while the chain length is bounded). tier-5 primes are 33-64 bit, so the
+# chain is ~16 base-16 Horner blocks (~1853 tokens). Vocab: digits 0..15 then
+# PAD,BOS,MUL,MOD,EQ,COLON,STEP,EOS = base..base+7 (see tier5_modmul.make_vocab).
+# The decoded answer is a BASE-16 residue; we convert it to base-10 digits with
+# multiply-add only (no %, //, Barrett/Montgomery/CRT on the product) so it matches
+# the global output_base=10. Compliance is unchanged: the only modular reduction in
+# shipped code is the per-operand int(a)%p / int(b)%p done before the network runs.
+# ---------------------------------------------------------------------------
+
+
+def _make_vocab_base(base: int) -> dict:
+    """Base-B scratchpad vocab, matching training/tier5_modmul.make_vocab."""
+    PAD, BOS, MUL, MOD, EQ, COLON, STEP, EOS = (
+        base, base + 1, base + 2, base + 3, base + 4, base + 5, base + 6, base + 7)
+    return dict(PAD=PAD, BOS=BOS, MUL=MUL, MOD=MOD, EQ=EQ, COLON=COLON, STEP=STEP,
+                EOS=EOS, VOCAB=base + 8,
+                SPECIALS={PAD, BOS, MUL, MOD, EQ, COLON, STEP, EOS})
+
+
+def _digits_base_msb(n: int, base: int) -> list[int]:
+    if n == 0:
+        return [0]
+    s = []
+    while n > 0:
+        s.append(n % base)
+        n //= base
+    return s[::-1]
+
+
+def _base_to_int(ds: list[int], base: int) -> int:
+    v = 0
+    for d in ds:
+        v = v * base + d        # multiply-add only; no %/// on the product
+    return v
+
+
+@torch.no_grad()
+def _modmul_decode_base(model, cfg, xyp, device, base, chunk=64):
+    """Greedy-decode (x*y) mod p in numeric base ``base`` for each (x, y, p) with
+    x, y already in [0, p). Returns base-10 digit-lists (MSB-first), or [0] if
+    unparseable. Mirrors _modmul_decode but base-parametrized; the final base-``base``
+    residue is re-expressed in base 10 via multiply-add (compliant)."""
+    V = _make_vocab_base(base)
+    PAD, EOS, COLON = V["PAD"], V["EOS"], V["COLON"]
+    max_len, abmax = cfg["max_len"], cfg["abacus_max"]
+    specials = torch.tensor(sorted(V["SPECIALS"]), device=device)
+    out: list[list[int] | None] = [None] * len(xyp)
+
+    groups = defaultdict(list)
+    prompts = []
+    for i, (x, y, p) in enumerate(xyp):
+        xd, yd, pd = (_digits_base_msb(x, base), _digits_base_msb(y, base),
+                      _digits_base_msb(p, base))
+        toks = [V["BOS"]] + xd + [V["MUL"]] + yd + [V["MOD"]] + pd + [V["EQ"]]
+        abac = ([0] + list(range(len(xd))) + [0] + list(range(len(yd)))
+                + [0] + list(range(len(pd))) + [0])
+        groups[len(toks)].append(i)
+        prompts.append((toks, abac))
+
+    for L, idxs in groups.items():
+        for s in range(0, len(idxs), chunk):
+            sub = idxs[s:s + chunk]
+            g = len(sub)
+            toks = torch.tensor([prompts[i][0] for i in sub], dtype=torch.long, device=device)
+            abac = torch.tensor([prompts[i][1] for i in sub], dtype=torch.long, device=device)
+            seg = torch.zeros(g, dtype=torch.long, device=device)
+            done = torch.zeros(g, dtype=torch.bool, device=device)
+            gen = [[] for _ in range(g)]
+            while toks.shape[1] < max_len and not bool(done.all()):
+                nxt = model(toks, abac)[:, -1].argmax(-1)
+                nxt = torch.where(done, torch.full_like(nxt, PAD), nxt)
+                is_sp = (nxt.unsqueeze(1) == specials).any(1)
+                new_abac = torch.where(is_sp, torch.zeros_like(seg),
+                                       torch.clamp(seg, max=abmax - 1))
+                seg = torch.where(is_sp, torch.zeros_like(seg), seg + 1)
+                nc, dc = nxt.tolist(), done.tolist()
+                for j in range(g):
+                    if not dc[j] and nc[j] != EOS and nc[j] != PAD:
+                        gen[j].append(nc[j])
+                toks = torch.cat([toks, nxt.unsqueeze(1)], dim=1)
+                abac = torch.cat([abac, new_abac.unsqueeze(1)], dim=1)
+                done = done | (nxt == EOS)
+            for j, i in enumerate(sub):
+                gj = gen[j]
+                if COLON in gj:
+                    k = len(gj) - 1 - gj[::-1].index(COLON)
+                    ans = [d for d in gj[k + 1:] if d < base]
+                    out[i] = int_to_decimal_digits(_base_to_int(ans, base)) if ans else [0]
+                else:
+                    out[i] = [0]
+            del toks, abac, seg, done, gen
+            # Only MPS needs the cache drop (it never frees mid-run and OOMs on
+            # long chains). On CUDA, empty_cache() synchronizes + forces the
+            # allocator to re-acquire buffers every chunk -- a ~5x slowdown over
+            # the ~1853-step base-16 decode -- and tier-5 peak is <3 GB anyway.
+            if device.type == "mps":
+                torch.mps.empty_cache()
+    return [o if o is not None else [0] for o in out]
+
+
 # ---------------------------------------------------------------------------
 # Submission entry class
 # ---------------------------------------------------------------------------
@@ -339,6 +449,8 @@ class EBMModMul(ModularMultiplicationModel):
         self.mm_cfg = None
         self.mm4 = None         # tier-4 modmul scratchpad
         self.mm4_cfg = None
+        self.mm5 = None         # tier-5 base-16 modmul scratchpad
+        self.mm5_cfg = None
 
     def load(self, model_dir: str) -> None:
         if torch.cuda.is_available():
@@ -375,6 +487,17 @@ class EBMModMul(ModularMultiplicationModel):
             ).to(self.device)
             self.mm4.load_state_dict(ckpt["tier4"]["state_dict"])
             self.mm4.eval()
+        # Tier 5: base-16 scratchpad, trained on primes in [2**33, 2**64).
+        if "tier5" in ckpt:
+            c5 = ckpt["tier5"]["config"]
+            self.mm5_cfg = c5
+            self.mm5 = AbacusDecoder(
+                max_len=c5["max_len"], abacus_max=c5["abacus_max"], d_model=c5["d_model"],
+                nhead=c5["nhead"], num_layers=c5["layers"], dim_ff=c5["dim_ff"],
+                vocab=c5["base"] + 8,
+            ).to(self.device)
+            self.mm5.load_state_dict(ckpt["tier5"]["state_dict"])
+            self.mm5.eval()
 
     # Per-argument identity preprocessing (each hook sees only its own argument).
     def preprocess_a(self, a): return a
@@ -392,6 +515,7 @@ class EBMModMul(ModularMultiplicationModel):
     TIER3_LO = 512
     TIER3_HI = 65536
     TIER4_HI = 2 ** 32
+    TIER5_HI = 2 ** 64
 
     @torch.no_grad()
     def predict_digits_batch(self, inputs):
@@ -399,16 +523,22 @@ class EBMModMul(ModularMultiplicationModel):
         x_rows, y_rows, p_rows, p_ints, idx = [], [], [], [], []   # tiers 1-2
         mm_items, mm_idx = [], []                                  # tier 3
         mm4_items, mm4_idx = [], []                                # tier 4
+        mm5_items, mm5_idx = [], []                                # tier 5
 
         for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
             p = int(p_enc)
             # Out of regime (residues don't fit the trained range): honest 0.
-            if p >= self.TIER4_HI:
+            if p >= self.TIER5_HI:
                 out[i] = [0]
                 continue
             a_red = int(a_enc) % p          # per-operand reduction (allowed)
             b_red = int(b_enc) % p
-            if p >= self.TIER3_HI:
+            if p >= self.TIER4_HI:
+                if self.mm5 is not None:
+                    mm5_items.append((a_red, b_red, p)); mm5_idx.append(i)
+                else:
+                    out[i] = [0]
+            elif p >= self.TIER3_HI:
                 if self.mm4 is not None:
                     mm4_items.append((a_red, b_red, p)); mm4_idx.append(i)
                 else:
@@ -445,6 +575,22 @@ class EBMModMul(ModularMultiplicationModel):
             for j, i in enumerate(mm4_idx):
                 out[i] = res[j]
 
+        if mm5_items:
+            # Tier-5 base-16 chains are ~1853 tokens. The model was trained under
+            # bf16 autocast, and decoding in bf16 (not fp32) is what keeps the long
+            # attention both fast (~125s/100 vs ~470-660s in fp32) and within memory
+            # (<3 GB vs a fp32 OOM at 1853-length attention). Match training precision.
+            if self.device.type == "cuda":
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                    res = _modmul_decode_base(self.mm5, self.mm5_cfg, mm5_items,
+                                              self.device, base=self.mm5_cfg["base"],
+                                              chunk=64)
+            else:
+                res = _modmul_decode_base(self.mm5, self.mm5_cfg, mm5_items,
+                                          self.device, base=self.mm5_cfg["base"], chunk=64)
+            for j, i in enumerate(mm5_idx):
+                out[i] = res[j]
+
         return [o if o is not None else [0] for o in out]
 
     def max_batch_size(self) -> int:

weights.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6d68b8f3ce76c7421b381a5c2465218d306323231386600f6ed643b889225ca6
-size 141288197
+oid sha256:165d028f8af0a9efb45b755a20e855701b3e6594bcbb69b63cdc356544db8303
+size 271364937