evaluation/evaluate_arithmark3.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ARITHMARK-3.0 RUNNER — INFERNO-EXP-1.7 / Deep400M
=============================================================
Runs AxiomicLabs/Arithmark-3.0 on the supplied Deep400M trainer's raw
checkpoint (best.pt / last.pt / epoch_N.pt), using the pinned official
bencharithmark-3.py scoring code. Evaluation only; no training state is resumed.

Expected model: 393,610,240 parameters with a 32,064-token vocabulary;
32 layers, width 1,024, 16 query / 4 KV heads, head dimension 64,
FFN inner width 2,816, canonical same-token SwiGLU and tied embeddings.
Architecture is read from checkpoint config/tensor shapes and checked below.

Faithfulness:
  • Downloads arithmark-3.jsonl and bencharithmark-3.py at a pinned revision,
    prints the dataset SHA-256, and calls their load_arithmark_3() and
    evaluate_arithmark_3() unchanged: separate context/continuation
    tokenization, no BOS/start token, context truncation at 1,024, right-padded
    batches of 8 examples (32 sequences; official default is 32 examples), raw summed and per-token
    length-normalized log-likelihood, acc and acc_norm (primary), per-category
    breakdown.
  • Uses the explicit FAST training tokenizer to reproduce the released score.
    This differs from the official loader's slow-tokenizer preference. Under
    Transformers 4.57.6, fast gives 85.70% acc_norm; slow gives 85.80%.
    Pad = eos when missing, TF32 matmuls on (as their main() sets), and for
    bfloat16 the weights themselves are cast like from_pretrained(dtype=bf16).
  • The model is an exact port of the GQA trainer's forward (strict state-dict load):
    fused qkv_proj, full-head half-split RoPE, SDPA enable_gqa, packed SwiGLU, tied embeddings.
EVAL_DTYPE: "float32" is the batch-size-stable reference most model cards
report; "bfloat16" is the script's own default. Report which one you used.

Hardcoded constants, no argparse. Colab: paste into one cell.
"""
import os
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ["HF_HUB_DISABLE_XET"] = "1"

import gc
import sys
import json
import math
import time
import hashlib
import importlib.util
from dataclasses import dataclass, asdict
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

# =============================================================================
# CONFIGURATION
# =============================================================================
# Same checkpoint selected by the supplied Deep400M trainer.
# Edit CHECKPOINT_DIR / CHECKPOINT_FILE to test another saved snapshot.
CHECKPOINT_DIR  = "/content/drive/MyDrive/Mathematician_L4_Experiments/gqa_deep400m_slm_v3_r1_anygpu/continue_r2_ctx4096"
CHECKPOINT_FILE = "32_indx.pt"      # best.pt / last.pt / epoch_N.pt; optimizer state is not needed
TOKENIZER_DIR   = ""            # "" → CHECKPOINT_DIR/tokenizer; use the tokenizer saved by training
MODEL_NAME      = "Inferno-Exp-1.7-Deep400M"
MOUNT_DRIVE_IF_COLAB = True

BENCH_REPO      = "AxiomicLabs/Arithmark-3.0"
BENCH_REVISION  = "6f6e59dd9b7e2c63455f7af7f838f9ecc3d0a746"   # current main; "main" = latest
BENCH_DIR       = "/content/arithmark3"

EVAL_DTYPE      = "float32"     # "float32" (reference) or "bfloat16" (script default)
BATCH_SIZE      = 8             # lower-memory default for 393.6M; official default is 32
                               # decrease to 4 / 2 / 1 if GPU memory is insufficient
MAX_CONTEXT     = 1024          # official default
SAVE_PREDICTIONS = True         # per-question scores in the JSON (useful for error analysis)
RESULTS_DIR     = "/content/arithmark3_results"

# Fail clearly if the selected checkpoint is from the old 76M run.
EXPECTED_ARCH = dict(d_model=1024, n_layers=32, n_heads=16, n_kv_heads=4,
                     head_dim=64, ffn_inner=2816, gate_shift_start=-1)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# =============================================================================
# OFFICIAL BENCHMARK CODE + DATA
# =============================================================================
def load_official_bench():
    from huggingface_hub import hf_hub_download
    folder = Path(BENCH_DIR); folder.mkdir(parents=True, exist_ok=True)
    data = Path(hf_hub_download(BENCH_REPO, "arithmark-3.jsonl", repo_type="dataset",
                                revision=BENCH_REVISION, local_dir=str(folder)))
    script = Path(hf_hub_download(BENCH_REPO, "bencharithmark-3.py", repo_type="dataset",
                                  revision=BENCH_REVISION, local_dir=str(folder)))
    spec = importlib.util.spec_from_file_location("axiomic_bencharithmark3", script)
    bench = importlib.util.module_from_spec(spec); sys.modules[spec.name] = bench
    spec.loader.exec_module(bench)
    return bench, data


# =============================================================================
# MODEL — exact port of the GQA trainers' forward (full-sequence)
# =============================================================================
@dataclass
class Arch:
    vocab_size: int
    d_model: int
    n_layers: int
    n_heads: int
    n_kv_heads: int
    head_dim: int
    ffn_inner: int
    gate_shift_start: int      # -1 = no gate shift (vanilla); else first shifted layer (0-based)
    context_len: int
    rope_base: float
    source_experiment: str

    @property
    def label(self):
        shift = "vanilla SwiGLU" if self.gate_shift_start < 0 else f"gate shift from layer {self.gate_shift_start}"
        return (f"GQA {self.n_layers}L×{self.d_model}d, {self.n_heads}Q/{self.n_kv_heads}KV heads × {self.head_dim}, "
                f"FFN {self.ffn_inner}, {shift}")


def arch_from_checkpoint(sd, config):
    """Head counts come from the stored config; every width is cross-checked against tensor shapes."""
    config = config or {}
    if "n_heads" not in config or "n_kv_heads" not in config:
        raise ValueError("Checkpoint has no n_heads / n_kv_heads in its config; is this a GQA trainer checkpoint?")
    d_model = sd["embed.weight"].shape[1]
    H, K = config["n_heads"], config["n_kv_heads"]
    if H <= 0 or K <= 0 or H % K:
        raise ValueError("Query heads must be a positive multiple of KV heads")
    q_width = sd["blocks.0.attn.o_proj.weight"].shape[1]
    if q_width % H:
        raise ValueError(f"o_proj input width {q_width} is not divisible by n_heads={H}")
    head_dim = q_width // H
    qkv_rows = sd["blocks.0.attn.qkv_proj.weight"].shape[0]
    if qkv_rows != (H + 2 * K) * head_dim:
        raise ValueError(f"qkv_proj has {qkv_rows} rows; expected (H+2K)*head_dim = {(H + 2 * K) * head_dim}")
    if head_dim % 2:
        raise ValueError("RoPE head dimension must be even")
    ffn_inner = sd["blocks.0.ffn.down.weight"].shape[1]
    if sd["blocks.0.ffn.project.weight"].shape[0] != 2 * ffn_inner:
        raise ValueError("FFN project/down widths disagree")
    variant = config.get("ffn_variant", "canonical_swiglu")
    if variant == "gate_shift_swiglu":
        shift = int(config.get("gate_shift_start", 1))
    elif variant == "canonical_swiglu":
        shift = -1
    else:
        raise ValueError(f"Unsupported ffn_variant {variant!r} for this runner")
    arch = Arch(vocab_size=sd["embed.weight"].shape[0], d_model=d_model,
                n_layers=1 + max(int(k.split(".")[1]) for k in sd if k.startswith("blocks.")),
                n_heads=H, n_kv_heads=K, head_dim=head_dim, ffn_inner=ffn_inner, gate_shift_start=shift,
                context_len=config.get("seq_len", 2048), rope_base=config.get("rope_base", 10000.),
                source_experiment=str(config.get("experiment_version")))
    for key in ("d_model", "n_layers", "ffn_inner"):
        if key in config and config[key] != getattr(arch, key):
            print(f"⚠️  {key}: config={config[key]} but tensor shapes say {getattr(arch, key)} → using shapes")
    mismatches = {key: (getattr(arch, key), expected)
                  for key, expected in EXPECTED_ARCH.items() if getattr(arch, key) != expected}
    if mismatches:
        details = ", ".join(f"{key}={actual} (expected {expected})"
                            for key, (actual, expected) in mismatches.items())
        raise ValueError(f"This runner expects the supplied Deep400M architecture: {details}")
    return arch


class RMSNorm(nn.RMSNorm):
    def forward(self, x):
        return F.rms_norm(x.float(), self.normalized_shape, self.weight, self.eps).to(x.dtype)


def rope_tables(length, dim, base):
    inv = 1. / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    freqs = torch.outer(torch.arange(length, dtype=torch.float32), inv)
    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    # Half-split convention over the FULL head dimension; FP32 phase math for AMP inputs.
    work = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x
    a, b = work.chunk(2, -1)
    return torch.cat((a * cos - b * sin, a * sin + b * cos), -1).to(x.dtype)


class GQAttention(nn.Module):
    def __init__(self, a):
        super().__init__()
        self.H = a.n_heads; self.K = a.n_kv_heads; self.hd = a.head_dim
        self.q_width = self.H * self.hd; self.kv_width = self.K * self.hd
        self.qkv_proj = nn.Linear(a.d_model, self.q_width + 2 * self.kv_width, bias=False)
        self.o_proj = nn.Linear(self.q_width, a.d_model, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        q, k, v = self.qkv_proj(x).split((self.q_width, self.kv_width, self.kv_width), -1)
        q = q.view(B, T, self.H, self.hd).transpose(1, 2)
        k = k.view(B, T, self.K, self.hd).transpose(1, 2)
        v = v.view(B, T, self.K, self.hd).transpose(1, 2)
        q = apply_rope(q, cos, sin); k = apply_rope(k, cos, sin)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0., enable_gqa=self.H != self.K)
        return self.o_proj(y.transpose(1, 2).reshape(B, T, self.q_width))


class FeedForward(nn.Module):
    def __init__(self, a, index):
        super().__init__()
        self.shift = a.gate_shift_start >= 0 and index >= a.gate_shift_start
        self.project = nn.Linear(a.d_model, 2 * a.ffn_inner, bias=False)
        self.down = nn.Linear(a.ffn_inner, a.d_model, bias=False)

    def forward(self, n):
        g, u = self.project(n).chunk(2, -1)
        if self.shift:
            g = torch.cat((g[:, :1], g[:, :-1]), 1)        # gate reads token t-1 (gate-shift trainer rule)
        return self.down(F.silu(g) * u)


class Block(nn.Module):
    def __init__(self, a, index):
        super().__init__()
        self.attn_norm = RMSNorm(a.d_model); self.attn = GQAttention(a)
        self.ffn_norm = RMSNorm(a.d_model); self.ffn = FeedForward(a, index)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.attn_norm(x), cos, sin)
        return x + self.ffn(self.ffn_norm(x))


class MathematicianGQA(nn.Module):
    def __init__(self, a):
        super().__init__()
        self.arch = a
        self.embed = nn.Embedding(a.vocab_size, a.d_model); self.emb_drop = nn.Dropout(0.)
        self.blocks = nn.ModuleList([Block(a, i) for i in range(a.n_layers)])
        self.ln_f = RMSNorm(a.d_model)
        self.lm_head = nn.Linear(a.d_model, a.vocab_size, bias=False); self.lm_head.weight = self.embed.weight
        cos, sin = rope_tables(a.context_len + 1, a.head_dim, a.rope_base)
        self.register_buffer("rope_cos", cos, persistent=False); self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, ids):
        if ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        T = ids.shape[1]
        if not 0 < T <= self.arch.context_len:
            raise ValueError(f"Sequence length {T} must be in [1, {self.arch.context_len}]")
        cos = self.rope_cos[:T]; sin = self.rope_sin[:T]
        x = self.emb_drop(self.embed(ids))
        for block in self.blocks:
            x = block(x, cos, sin)
        return F.linear(self.ln_f(x), self.embed.weight).float()


class _Out:
    def __init__(self, logits):
        self.logits = logits


class HFStyleAdapter(nn.Module):
    """What evaluate_arithmark_3 calls: model(tokens, attention_mask=..., use_cache=False).logits.

    Their batches are right-padded and only positions before each sequence's end are
    scored, so a plain causal forward is exact; attention_mask is accepted and unused.
    Autocast (bf16 mode) is applied by their evaluate function, as for any HF model.
    """
    def __init__(self, inner):
        super().__init__()
        self.inner = inner

    def forward(self, input_ids=None, attention_mask=None, use_cache=None, **kwargs):
        return _Out(self.inner(input_ids))


# =============================================================================
# LOADING
# =============================================================================
def maybe_mount_drive():
    if not (MOUNT_DRIVE_IF_COLAB and CHECKPOINT_DIR.startswith("/content/drive")):
        return
    try:
        import google.colab  # noqa: F401
    except ImportError:
        return
    if not os.path.ismount("/content/drive"):
        from google.colab import drive
        drive.mount("/content/drive")


def load_model(dtype):
    path = os.path.join(CHECKPOINT_DIR, CHECKPOINT_FILE)
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Checkpoint not found: {path}. Set CHECKPOINT_DIR / CHECKPOINT_FILE.")
    # A full last.pt/epoch_N.pt may include several GB of AdamW state; only model
    # weights are loaded into the evaluator, and the saved payload is then freed.
    saved = torch.load(path, map_location="cpu", weights_only=True)
    sd = saved["model"] if isinstance(saved, dict) and "model" in saved else saved
    sd = {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
    meta = {k: v for k, v in saved.items() if k in ("step", "epoch", "best_nll", "eos_token_id")} \
        if isinstance(saved, dict) else {}
    arch = arch_from_checkpoint(sd, saved.get("config") if isinstance(saved, dict) else None)
    model = MathematicianGQA(arch)
    model.load_state_dict(sd, strict=True)
    del saved, sd; gc.collect()
    if dtype != torch.float32:
        # Same effect as from_pretrained(dtype=...): parameters in reduced precision.
        # RoPE tables stay fp32 buffers (apply_rope does its phase math in fp32 anyway).
        for p in model.parameters():
            p.data = p.data.to(dtype)
    model.to(DEVICE).eval()
    return model, arch, meta


def load_tokenizer(arch, meta):
    """Explicit fast training tokenizer reproduces the released 85.70% result.

    Pin transformers==4.57.6 for this reproduction. A genuine slow LlamaTokenizer
    produces a separate 85.80% normalized / 86.00% raw reference on these weights.
    """
    from transformers import AutoTokenizer
    folder = TOKENIZER_DIR or os.path.join(CHECKPOINT_DIR, "tokenizer")
    if not os.path.isdir(folder):
        raise FileNotFoundError(f"Training tokenizer folder not found: {folder}")
    tok = AutoTokenizer.from_pretrained(folder, use_fast=True, local_files_only=True)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    if len(tok) != arch.vocab_size:
        raise ValueError(f"Tokenizer has {len(tok)} tokens; checkpoint expects {arch.vocab_size}")
    if "eos_token_id" in meta and tok.eos_token_id != meta["eos_token_id"]:
        raise ValueError("Tokenizer EOS differs from the selected checkpoint")
    return tok, f"{type(tok).__name__} (fast training tokenizer, release reference)"


# =============================================================================
# MAIN — mirrors the official main() for one model
# =============================================================================
def main():
    maybe_mount_drive()
    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[EVAL_DTYPE]
    torch.set_float32_matmul_precision("high")          # as their main() does
    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    bench, data_path = load_official_bench()
    examples = bench.load_arithmark_3(data_path)
    dataset_sha256 = hashlib.sha256(data_path.read_bytes()).hexdigest()
    print(f"Loaded {len(examples)} ArithMark 3.0 examples @ {BENCH_REVISION[:8]}")
    print(f"ArithMark 3.0 dataset SHA-256: {dataset_sha256}")

    model, arch, meta = load_model(dtype)
    tok, tok_kind = load_tokenizer(arch, meta)
    adapter = HFStyleAdapter(model).eval()
    total_params = sum(p.numel() for p in model.parameters())
    epoch = f" · epoch {meta['epoch']}" if meta.get("epoch") is not None else ""
    print(f"{MODEL_NAME}: {total_params:,} parameters ({dtype}) · {arch.label} · {arch.source_experiment} · "
          f"step {meta.get('step')}{epoch} · tokenizer {tok_kind} · {DEVICE}")

    print(f"Evaluation: batch_size={BATCH_SIZE} examples ({4 * BATCH_SIZE} choices), "
          f"max_context={MAX_CONTEXT}, dtype={EVAL_DTYPE}")
    with torch.inference_mode():
        result = bench.evaluate_arithmark_3(adapter, tok, DEVICE, dtype, examples,
                                            BATCH_SIZE, MAX_CONTEXT, SAVE_PREDICTIONS)
    result["primary_metric"] = "acc_norm"
    result["primary_acc"] = result["acc_norm"]

    print(f"\n{'=' * 68}\n  {MODEL_NAME} ({total_params:,} params) RESULTS\n{'=' * 68}")
    print(f"  Raw continuation accuracy     {result['acc']:>8.2f}%")
    print(f"  Length-normalized accuracy    {result['acc_norm']:>8.2f}%")
    print(f"  Primary (acc_norm){result['acc_norm']:>16.2f}%")
    se = math.sqrt(result['acc_norm'] / 100 * (1 - result['acc_norm'] / 100) / result['total']) * 100
    print(f"  Standard error (acc_norm)     ±{se:.2f} pts   · chance 25.00%")
    print('=' * 68)

    out_dir = Path(RESULTS_DIR); out_dir.mkdir(parents=True, exist_ok=True)
    out_file = out_dir / f"{MODEL_NAME}_arithmark-3_results.json"
    with out_file.open("w", encoding="utf-8") as handle:
        json.dump({
            "model": MODEL_NAME,
            "tokenizer": f"{TOKENIZER_DIR or os.path.join(CHECKPOINT_DIR, 'tokenizer')} [{tok_kind}]",
            "params": total_params,
            "dataset": {"path": str(data_path), "sha256": dataset_sha256, "items": len(examples),
                        "revision": BENCH_REVISION},
            "runtime": {"device": DEVICE, "dtype": str(dtype), "batch_size": BATCH_SIZE,
                        "max_context": MAX_CONTEXT, "attention_implementation": "sdpa enable_gqa (custom GQA port)",
                        "compiled": False, "torch_version": torch.__version__},
            "checkpoint": {"path": os.path.join(CHECKPOINT_DIR, CHECKPOINT_FILE), "step": meta.get("step"),
                           "epoch": meta.get("epoch"),
                           "architecture": asdict(arch)},
            "primary_metric": "acc_norm",
            "results": {"arithmark-3": result},
        }, handle, indent=2)
    print(f"Results saved to {out_file}")
    return result


if __name__ == "__main__":
    main()