verify_paste.py

#!/usr/bin/env python3
"""
verify_paste.py — Verify LayerCake domain paste is bit-exact and function-preserving.

No GPU required. No data required. No checkpoints required. Runs in ~10 seconds.

What this proves:
  1. Weight identity      — pasted domain module tensors are bit-exact copies
  2. Forward-pass identity — same core + pasted domain → identical logits (max_diff = 0.0)
  3. Generation identity  — 50-token autoregressive sequence is token-for-token identical
  4. Cross-size portability — domain function is preserved across 48M and 150M model sizes

Usage:
    python verify_paste.py
"""

import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

import torch
from model import LayerCakeLMFixedABI

# ── Config ──────────────────────────────────────────────────────────────────
VOCAB    = 1000
D_ABI    = 512
SEQ_LEN  = 32
BATCH    = 4
SEED     = 42
DOMAINS  = ["chess", "python"]
SMALL    = dict(d_model=512, n_layers=4, n_heads=8)
LARGE    = dict(d_model=768, n_layers=6, n_heads=12)

def _build(d_model, n_layers, n_heads):
    return LayerCakeLMFixedABI(
        vocab_size=VOCAB, d_model=d_model, d_abi=D_ABI,
        n_core_layers=n_layers, n_heads=n_heads, d_ff=d_model*4,
        domain_names=DOMAINS, max_seq_len=SEQ_LEN,
        use_router=False, domain_module_type="lite",
    )

def _paste(src, tgt, domain):
    ss, ts = src.state_dict(), tgt.state_dict()
    prefix = f"domain_modules.{domain}."
    for k, v in ss.items():
        if k.startswith(prefix):
            ts[k] = v.clone()
    tgt.load_state_dict(ts)

def _header(title):
    print(f"\n  [{title}]")

def _ok(msg):
    print(f"    ✓  {msg}")

def _fail(msg):
    print(f"    ✗  {msg}")
    sys.exit(1)

# ── Test 1: Weight identity ──────────────────────────────────────────────────
def test_weight_identity():
    _header("weight identity — 9 tensors, cross-size 48M → 150M")
    torch.manual_seed(SEED)
    small = _build(**SMALL)
    torch.manual_seed(SEED + 1)
    large = _build(**LARGE)
    _paste(small, large, "chess")
    ss = {k: v for k, v in small.state_dict().items() if k.startswith("domain_modules.chess.")}
    ls = {k: v for k, v in large.state_dict().items() if k.startswith("domain_modules.chess.")}
    for k in ss:
        d = (ss[k] - ls[k]).abs().max().item()
        if d != 0.0:
            _fail(f"tensor {k} max_diff={d}")
    _ok(f"{len(ss)} tensors — max_diff = 0.0 (bit-exact)")

# ── Test 2: Forward-pass identity ────────────────────────────────────────────
def test_forward_pass():
    _header("forward-pass identity — same core + pasted domain")
    torch.manual_seed(SEED)
    model_a = _build(**SMALL)
    with torch.no_grad():
        for n, p in model_a.named_parameters():
            if n.startswith("domain_modules.chess."):
                p.add_(torch.randn_like(p) * 0.5)

    torch.manual_seed(SEED + 77)
    model_ap = _build(**SMALL)
    sa, sap = model_a.state_dict(), model_ap.state_dict()
    for k in sa:
        if not k.startswith("domain_modules."):
            sap[k] = sa[k].clone()
    model_ap.load_state_dict(sap)
    _paste(model_a, model_ap, "chess")

    torch.manual_seed(999)
    x = torch.randint(0, VOCAB, (BATCH, SEQ_LEN))
    mask = torch.zeros(BATCH, len(DOMAINS))
    mask[:, DOMAINS.index("chess")] = 1.0

    model_a.eval(); model_ap.eval()
    with torch.no_grad():
        la, _ = model_a(x, domain_mask=mask)
        lap, _ = model_ap(x, domain_mask=mask)

    d = (la - lap).abs().max().item()
    if d != 0.0:
        _fail(f"logit max_diff = {d:.2e}")
    _ok(f"logit max_diff = 0.000000e+00 (bit-exact)")

# ── Test 3: Generation identity ──────────────────────────────────────────────
def test_generation():
    _header("generation identity — 50-token autoregressive sequence")
    torch.manual_seed(SEED)
    model_a = _build(**SMALL)
    with torch.no_grad():
        for n, p in model_a.named_parameters():
            if n.startswith("domain_modules.chess."):
                p.add_(torch.randn_like(p) * 0.5)

    torch.manual_seed(SEED + 111)
    model_ap = _build(**SMALL)
    sa, sap = model_a.state_dict(), model_ap.state_dict()
    for k in sa:
        if not k.startswith("domain_modules."):
            sap[k] = sa[k].clone()
    model_ap.load_state_dict(sap)
    _paste(model_a, model_ap, "chess")

    torch.manual_seed(999)
    prompt = torch.randint(0, VOCAB, (1, 8))
    mask = torch.zeros(1, len(DOMAINS))
    mask[:, DOMAINS.index("chess")] = 1.0

    model_a.eval(); model_ap.eval()
    toks_a, toks_ap = prompt.clone(), prompt.clone()

    with torch.no_grad():
        for step in range(50):
            la, _ = model_a(toks_a, domain_mask=mask)
            lap, _ = model_ap(toks_ap, domain_mask=mask)
            na = la[:, -1, :].argmax(-1, keepdim=True)
            nap = lap[:, -1, :].argmax(-1, keepdim=True)
            if na.item() != nap.item():
                _fail(f"sequences diverged at step {step+1}: {na.item()} vs {nap.item()}")
            toks_a  = torch.cat([toks_a,  na],  dim=1)
            toks_ap = torch.cat([toks_ap, nap], dim=1)

    gen = toks_a[0, 8:].tolist()
    _ok(f"50 tokens generated — zero divergence — e.g. {gen[:8]}…")

# ── Test 4: Cross-size domain function ───────────────────────────────────────
def test_cross_size_function():
    _header("cross-size function identity — domain module 48M → 150M")
    torch.manual_seed(SEED)
    small = _build(**SMALL)
    with torch.no_grad():
        for n, p in small.named_parameters():
            if n.startswith("domain_modules.chess."):
                p.add_(torch.randn_like(p) * 0.5)

    torch.manual_seed(SEED + 99)
    large = _build(**LARGE)
    _paste(small, large, "chess")

    torch.manual_seed(42)
    h = torch.randn(BATCH, SEQ_LEN, D_ABI)
    small.eval(); large.eval()
    with torch.no_grad():
        out_s = small.domain_modules["chess"](h)
        out_l = large.domain_modules["chess"](h)

    d = (out_s - out_l).abs().max().item()
    if d != 0.0:
        _fail(f"domain output max_diff = {d:.2e}")
    _ok(f"domain output max_diff = 0.000000e+00 (bit-exact, d_model 512→768)")


# ── Main ─────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print("=" * 56)
    print("  LayerCake — Paste Verification")
    print("  No GPU · No data · No checkpoints required")
    print("=" * 56)

    t0 = time.time()
    test_weight_identity()
    test_forward_pass()
    test_generation()
    test_cross_size_function()
    elapsed = time.time() - t0

    print(f"\n{'=' * 56}")
    print(f"  ALL CHECKS PASSED  ({elapsed:.1f}s)")
    print(f"  Paste is bit-exact and function-preserving.")
    print(f"  See CLAIMS.md and SKEPTICS.md for full scope.")
    print(f"{'=' * 56}\n")