File size: 6,448 Bytes
63f41a6 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | #!/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("<Q", f.read(8))
if n > 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)
|