#!/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()