File size: 3,002 Bytes
5e9c921 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """
smoke_gpu.py -- Step 2 self-checks that need the real model on a GPU.
Single model load covers:
A. target-linear enumeration: all=224, attn=128, mlp=96 (32 blocks x 7).
B. lm_head hard defense passes on the real dense model.
C. 3-item REAL-eval schema check: eval path dumps valid per-item JSONL that
mcnemar.py can read (item_id unique, fields complete). Scores NOT inspected.
"""
import os
import sys
import json
import tempfile
import importlib.util
import torch
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import common as C
from analysis import mcnemar
spec = importlib.util.spec_from_file_location("tm_eval", os.path.join(HERE, "eval", "eval.py"))
E = importlib.util.module_from_spec(spec)
spec.loader.exec_module(E)
ok = True
# ── A. counts ─────────────────────────────────────────────────────────────────
print("[A] loading dense model...")
model, tok = C.load_model()
n_all = C.count_target_linears(model, "all")
n_attn = C.count_target_linears(model, "attn")
n_mlp = C.count_target_linears(model, "mlp")
a_ok = (n_all == 224 and n_attn == 128 and n_mlp == 96)
ok &= a_ok
print(f"[A] target linears all={n_all} attn={n_attn} mlp={n_mlp} "
f"(expect 224/128/96) -> {'PASS' if a_ok else 'FAIL'}")
# ── B. head guard on real model ───────────────────────────────────────────────
try:
C.assert_head_dense(model)
head = C.get_output_head(model)
b_ok = True
print(f"[B] head is {type(head).__name__} (dense) -> PASS")
except Exception as e:
b_ok = False
print(f"[B] head guard FAILED on dense model: {e}")
ok &= b_ok
# ── C. 3-item real-eval schema dump (MMLU, ref arm) ───────────────────────────
from datasets import load_dataset
rows = load_dataset("cais/mmlu", "all", split="test")
rows = [rows[i] for i in range(3)]
fd, path = tempfile.mkstemp(suffix="_smoke_items.jsonl")
with os.fdopen(fd, "w") as f:
c, cn, tot = E._run_mcq(model, tok, f, "mmlu", rows, E._mmlu_task,
mc_num=32, batch_size=16, cfg=0.5,
shard=0, num_shards=1, limit=3)
read = mcnemar.read_items(path)
required = {"item_id", "prompt_hash", "correct", "correct_norm", "pred", "gold"}
unique = len(read) == 3
complete = all(required <= set(v) for v in read.values())
typed = all(isinstance(v["correct"], bool) for v in read.values())
c_ok = unique and complete and typed and tot == 3
ok &= c_ok
print(f"[C] 3-item MMLU dump: n={tot} unique={unique} complete={complete} "
f"bool={typed} -> {'PASS' if c_ok else 'FAIL'}")
print(f" sample record: {json.dumps(list(read.values())[0])}")
os.unlink(path)
print(f"\n=== GPU SMOKE {'ALL PASS' if ok else 'FAIL'} ===")
sys.exit(0 if ok else 1)
|