File size: 4,769 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 | #!/usr/bin/env python
"""CPU-only compatibility check of a HF snapshot against align's Qwen3 loader (no weights are loaded).
.venv/bin/python scripts/p2/check_qwen3_snapshot.py SNAPSHOT_DIR [--json]
1. config.json parses through `align.model.qwen3_checkpoint.read_qwen3_config` (it REFUSES every variant align does
not implement: sliding window, scaled RoPE, biases, tied head, non-silu, ...);
2. the safetensors index's `weight_map` keys equal EXACTLY the key set `hf_to_align` maps (no missing, no extra);
3. every shard exists and its header-declared tensor extents fit the file (a truncated download fails here);
4. every tensor's shape (read from the shard HEADER, 8 bytes + JSON, no tensor data) equals the shape the config implies;
5. dtypes are reported (bf16 expected).
Exit 0 compatible, 1 not. Memory: a few MB. Threads: 1.
"""
from __future__ import annotations
import argparse
import json
import os
import struct
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from align.model.checkpoint import INDEX_NAME # noqa: E402
from align.model.qwen3_checkpoint import _expected_shapes, _table, read_qwen3_config # noqa: E402
def read_header(path: str) -> dict:
"""The safetensors header dict, plus `__file_ok__`: whether the declared extents fit the file."""
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}: implausible safetensors 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["__file_ok__"] = 8 + n + end == size
h["__declared_bytes__"] = 8 + n + end
h["__file_bytes__"] = size
return h
def check(snapshot: str) -> dict:
snap = Path(snapshot)
rep: dict = {"snapshot": str(snap), "problems": []}
P = rep["problems"]
cfg = read_qwen3_config(str(snap))
rep["config"] = {k: getattr(cfg, k) for k in ("n_layers", "hidden_size", "intermediate_size", "n_heads", "n_kv_heads", "head_dim",
"vocab_size", "rope_theta", "rms_norm_eps", "max_position_embeddings", "tie_word_embeddings")}
want = set(_table(cfg))
shapes = _expected_shapes(cfg)
index = snap / INDEX_NAME
if not index.is_file():
P.append(f"no {INDEX_NAME}")
return rep
wm = json.loads(index.read_text())["weight_map"]
rep["index"] = {"tensors": len(wm), "shards": sorted(set(wm.values()))}
missing, extra = sorted(want - set(wm)), sorted(set(wm) - want)
if missing:
P.append(f"{len(missing)} key(s) hf_to_align needs are absent from the index: {missing[:5]}")
if extra:
P.append(f"{len(extra)} index key(s) hf_to_align cannot map: {extra[:5]}")
dtypes: dict = {}
seen = 0
for shard in rep["index"]["shards"]:
sp = snap / shard
if not sp.is_file():
P.append(f"shard {shard} named by the index is absent")
continue
h = read_header(str(sp))
if not h["__file_ok__"]:
P.append(f"shard {shard}: header declares {h['__declared_bytes__']} bytes, file has {h['__file_bytes__']} (truncated or padded)")
for k, v in h.items():
if k.startswith("__"):
continue
seen += 1
dtypes[v["dtype"]] = dtypes.get(v["dtype"], 0) + 1
if wm.get(k) != shard:
P.append(f"{k}: in {shard} but the index says {wm.get(k)}")
if k in shapes and tuple(v["shape"]) != shapes[k]:
P.append(f"{k}: shape {tuple(v['shape'])} != {shapes[k]} implied by the config")
rep["tensors_seen_in_shard_headers"], rep["dtypes"] = seen, dtypes
if seen != len(wm):
P.append(f"shard headers carry {seen} tensors, the index lists {len(wm)}")
rep["compatible"] = not P
return rep
def main(argv=None) -> int:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("snapshot")
ap.add_argument("--json", action="store_true")
a = ap.parse_args(argv)
try:
rep = check(a.snapshot)
except Exception as e: # noqa: BLE001 -- a refusal from the config parser is the answer
rep = {"snapshot": a.snapshot, "compatible": False, "problems": [f"{type(e).__name__}: {e}"]}
print(json.dumps(rep, indent=1, default=str) if a.json else
("COMPATIBLE" if rep.get("compatible") else "NOT COMPATIBLE") + "\n" + "\n".join(f" - {p}" for p in rep["problems"]))
return 0 if rep.get("compatible") else 1
if __name__ == "__main__":
sys.exit(main())
|