| |
| """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": |
| 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: |
| |
| |
| |
| 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) |
|
|
