#!/usr/bin/env python """Header-only completeness + config check of a HF safetensors snapshot (Llama / Qwen2 / Qwen3 dense). No weights loaded. python3 -I scripts/teachers/check_snapshot.py SNAPSHOT_DIR [--json] Checks: (1) config.json architecture is one this script knows and has no feature a vLLM-default loader would mishandle silently (reports rope_scaling / sliding_window / tie_word_embeddings instead of refusing); (2) the index `weight_map` names every tensor once and every shard it references exists; (3) each shard header parses and its declared extents equal the file size (a truncated download fails here); (4) a tensor held by a different shard than the index names is a WARNING (--strict-index: a problem); (5) each tensor shape equals the shape the config implies; (6) dtypes. Exit 0 = ok, 1 = problem. Memory: a few MB. Does not replace scripts/p2/check_qwen3_snapshot.py, which additionally proves the key set maps into align's loader. """ from __future__ import annotations import argparse import json import os import struct import sys from pathlib import Path def read_header(path: Path) -> dict: size = os.path.getsize(path) with open(path, "rb") as f: (n,) = struct.unpack(" 100 * 1024 * 1024 or 8 + n > size: raise ValueError(f"{path.name}: implausible header length {n}") h = json.loads(f.read(n)) end = max((v["data_offsets"][1] for k, v in h.items() if k != "__metadata__"), default=0) h["__ok__"] = 8 + n + end == size return h def expected_shapes(c: dict) -> dict[str, list[int]]: L, H, I, V = c["num_hidden_layers"], c["hidden_size"], c["intermediate_size"], c["vocab_size"] nh, nkv = c["num_attention_heads"], c.get("num_key_value_heads", c["num_attention_heads"]) hd = c.get("head_dim") or H // nh arch = c["architectures"][0] s = {"model.embed_tokens.weight": [V, H], "model.norm.weight": [H]} if not c.get("tie_word_embeddings", False): s["lm_head.weight"] = [V, H] for i in range(L): p = f"model.layers.{i}." s[p + "self_attn.q_proj.weight"] = [nh * hd, H] s[p + "self_attn.k_proj.weight"] = [nkv * hd, H] s[p + "self_attn.v_proj.weight"] = [nkv * hd, H] s[p + "self_attn.o_proj.weight"] = [H, nh * hd] s[p + "mlp.gate_proj.weight"] = [I, H] s[p + "mlp.up_proj.weight"] = [I, H] s[p + "mlp.down_proj.weight"] = [H, I] s[p + "input_layernorm.weight"] = [H] s[p + "post_attention_layernorm.weight"] = [H] if arch == "Qwen3ForCausalLM": s[p + "self_attn.q_norm.weight"] = [hd] s[p + "self_attn.k_norm.weight"] = [hd] if arch == "Qwen2ForCausalLM": # q/k/v biases s[p + "self_attn.q_proj.bias"] = [nh * hd] s[p + "self_attn.k_proj.bias"] = [nkv * hd] s[p + "self_attn.v_proj.bias"] = [nkv * hd] return s def check(snap: Path, strict_index: bool = False) -> dict: rep: dict = {"snapshot": str(snap), "problems": []} P = rep["problems"] c = json.loads((snap / "config.json").read_text()) arch = (c.get("architectures") or ["?"])[0] rep["config"] = {k: c.get(k) for k in ("architectures", "num_hidden_layers", "hidden_size", "intermediate_size", "num_attention_heads", "num_key_value_heads", "head_dim", "vocab_size", "max_position_embeddings", "rope_theta", "rope_scaling", "sliding_window", "tie_word_embeddings", "torch_dtype", "quantization_config")} if arch not in ("LlamaForCausalLM", "Qwen2ForCausalLM", "Qwen3ForCausalLM"): P.append(f"unknown architecture {arch}") return rep idx = snap / "model.safetensors.index.json" if not idx.is_file(): P.append("no model.safetensors.index.json") return rep wm = json.loads(idx.read_text())["weight_map"] want = expected_shapes(c) missing, extra = sorted(set(want) - set(wm)), sorted(set(wm) - set(want)) rep["tensors"] = {"index": len(wm), "expected": len(want)} if missing: P.append(f"{len(missing)} tensors missing from index, e.g. {missing[:3]}") if extra: P.append(f"{len(extra)} unexpected tensors in index, e.g. {extra[:3]}") dtypes: dict[str, int] = {} total = 0 seen: set[str] = set() mismatched: list[str] = [] for shard in sorted(set(wm.values())): f = snap / shard if not f.is_file(): P.append(f"missing shard {shard}") continue total += f.stat().st_size h = read_header(f) if not h["__ok__"]: P.append(f"shard {shard} truncated or padded (header extents != file size)") for k, v in h.items(): if k in ("__metadata__", "__ok__"): continue seen.add(k) dtypes[v["dtype"]] = dtypes.get(v["dtype"], 0) + 1 if wm.get(k) != shard: mismatched.append(k) if k in want and v["shape"] != want[k]: P.append(f"{k}: shape {v['shape']} != expected {want[k]}") if set(wm) - seen: P.append(f"{len(set(wm) - seen)} indexed tensors absent from shard headers") rep["dtypes"], rep["bytes"] = dtypes, total if mismatched: # every tensor exists exactly once, but the index points some at a different shard than the one holding them. A loader that # opens every shard (vLLM, transformers) is unaffected; one that follows weight_map per tensor would fail. Upstream defect # in Goedel-Prover-V2-32B (index byte-identical to the hub's). Warning, or a problem with --strict-index. msg = f"{len(mismatched)} tensors are in a different shard than the index says, e.g. {mismatched[:2]}" (P if strict_index else rep.setdefault("warnings", [])).append(msg) rep["ok"] = not P return rep if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("snapshot") ap.add_argument("--json", action="store_true") ap.add_argument("--strict-index", action="store_true", help="treat index-vs-shard assignment mismatches as problems") a = ap.parse_args() r = check(Path(a.snapshot), a.strict_index) print(json.dumps(r, indent=1) if a.json else json.dumps({k: v for k, v in r.items() if k != "config"}, indent=1)) print("config:", json.dumps(r["config"])) sys.exit(0 if not r["problems"] else 1)