| |
| """ |
| in2_gate.py -- the EXACT layout gate for in2_gen.py output. No GPU, no CUDA. |
| |
| This is the only exact pre-flight check in the IN2 pilot. The step-0 loss check |
| was demoted to a diagnostic because its two hypotheses sit only ~0.22 nats apart |
| (no-induction 4.147 vs strong-induction 3.925, against a 3.901 plain baseline), |
| which cannot validate a layout. These assertions can. |
| |
| Checks, per arm dir: |
| 1 file size is exactly examples*T*2 bytes (headerless uint16) |
| 2 every token in [0, VOCAB) |
| 3 no answer span intrudes into the last PAD_TAIL offsets of its example |
| 4 planted spans (needle / probe+answer / decoy) never overlap |
| 5 RETRIEVAL IS ACTUALLY POSSIBLE: the tokens at the probe equal the needle's |
| key, and the tokens at the answer equal the needle's value. If this fails |
| the task is unsolvable and the whole experiment is void. |
| 6 batch alignment: re-derive x/y exactly as DataLoaderLite::next_batch does |
| (x=buf[i], y=buf[i+1], stride B*T*dp) and assert every example boundary |
| lands on a row boundary |
| 7 arm-specific distance invariants (b: all == span, c: all >= d_min) |
| |
| Cross-arm: |
| 8 arms B and C have BIT-IDENTICAL answer_start histograms -- the discipline |
| the whole B-vs-C contrast rests on |
| |
| Usage |
| python3 in2_gate.py --dirs /tmp/in2b /tmp/in2c --B 2 --dp 2 |
| """ |
| import argparse, array, json, os, sys |
|
|
| VOCAB = 50304 |
| fails = [] |
|
|
|
|
| def check(cond, msg): |
| print((" ok " if cond else " FAIL ") + msg) |
| if not cond: |
| fails.append(msg) |
|
|
|
|
| def load(d, split="train"): |
| js = [f for f in os.listdir(d) if f.endswith(".json") and f"_{split}_" in f] |
| if len(js) != 1: |
| raise SystemExit(f"{d}: expected exactly one {split} manifest, found {js}") |
| man = json.load(open(os.path.join(d, js[0]))) |
| toks = array.array("H") |
| for sp in man.get("shards", [man["bin"]]): |
| with open(sp, "rb") as f: |
| toks.fromfile(f, os.path.getsize(sp) // 2) |
| return man, toks |
|
|
|
|
| def gate_one(d, B, dp, split="train"): |
| man, toks = load(d, split) |
| T, n_ex, arm, span = man["T"], man["examples"], man["arm"], man["span"] |
| kl, vl, pad = man["key_len"], man["val_len"], man["pad_tail"] |
| print(f"\n=== arm {arm} T={T} examples={n_ex} K={man['K']} ===") |
|
|
| tot = sum(os.path.getsize(sp) for sp in man.get("shards", [man["bin"]])) |
| check(tot == n_ex * T * 2, |
| f"1 total size across {len(man.get('shards',[1]))} shard(s) == examples*T*2 " |
| f"({n_ex*T*2} bytes, got {tot})") |
| lo, hi = min(toks), max(toks) |
| check(lo >= 0 and hi < VOCAB, f"2 tokens in [0,{VOCAB}) (min {lo}, max {hi})") |
|
|
| bad_pad = bad_ovl = bad_ret = bad_dist = 0 |
| for ei, recs in enumerate(man["records"]): |
| base = ei * T |
| spans = [] |
| for r in recs: |
| a0, a1 = r["answer_start"], r["answer_start"] + vl |
| if a1 > T - pad: |
| bad_pad += 1 |
| spans.append((r["needle_start"], r["needle_start"] + span)) |
| spans.append((r["probe_start"], r["probe_start"] + span)) |
| |
| nk = toks[base + r["needle_start"]: base + r["needle_start"] + kl] |
| nv = toks[base + r["needle_start"] + kl: base + r["needle_start"] + span] |
| pk = toks[base + r["probe_start"]: base + r["probe_start"] + kl] |
| av = toks[base + a0: base + a1] |
| if nk != pk or nv != av: |
| bad_ret += 1 |
| if arm == "b" and r["distance"] != span: |
| bad_dist += 1 |
| if arm == "c" and r["distance"] < man["d_min"]: |
| bad_dist += 1 |
| spans.sort() |
| for (s0, e0), (s1, _) in zip(spans, spans[1:]): |
| if s1 < e0: |
| bad_ovl += 1 |
| check(bad_pad == 0, f"3 no answer within the last {pad} offsets ({bad_pad} bad)") |
| check(bad_ovl == 0, f"4 planted spans never overlap ({bad_ovl} bad)") |
| check(bad_ret == 0, f"5 RETRIEVAL POSSIBLE: probe==key and answer==value ({bad_ret} bad)") |
| check(bad_dist == 0, f"7 arm-{arm} distance invariant ({bad_dist} bad)") |
|
|
| |
| |
| |
| stride = B * T * dp |
| check(stride % T == 0, f"6a batch stride {stride} is a multiple of T={T}") |
| check((n_ex * T) % stride == 0, |
| f"6b total {n_ex*T} tokens is a whole number of batches (stride {stride})") |
| if len(toks) >= B * T + 1: |
| x0 = toks[0:B * T] |
| y0 = toks[1:B * T + 1] |
| rows_ok = all(x0[r * T] == toks[r * T] for r in range(B)) |
| check(rows_ok, f"6c row r starts at token r*T for all B={B} rows") |
| check(y0[T - 1] == toks[T], "6d y[T-1] targets the next example's first token " |
| "(why PAD_TAIL exists)") |
| return man |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--dirs", nargs="+", required=True) |
| ap.add_argument("--B", type=int, default=2) |
| ap.add_argument("--dp", type=int, default=2) |
| ap.add_argument("--split", default="train", choices=["train", "val"]) |
| args = ap.parse_args() |
|
|
| mans = {} |
| for d in args.dirs: |
| m = gate_one(d, args.B, args.dp, args.split) |
| mans[m["arm"]] = m |
|
|
| if "b" in mans and "c" in mans: |
| print("\n=== cross-arm ===") |
| hb, hc = mans["b"]["answer_start_hist"], mans["c"]["answer_start_hist"] |
| check(hb == hc, "8 arms B and C have BIT-IDENTICAL answer_start histograms " |
| f"({len(hb)} vs {len(hc)} distinct offsets)") |
| check(mans["b"]["seed"] == mans["c"]["seed"], "8b same seed") |
| check(mans["b"]["answer_tokens"] == mans["c"]["answer_tokens"], |
| "8c same answer-token count") |
|
|
| print(f"\n{'ALL PASS' if not fails else str(len(fails)) + ' FAILURE(S)'}") |
| return 1 if fails else 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
