| |
| """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 |
| from align.model.qwen3_checkpoint import _expected_shapes, _table, read_qwen3_config |
|
|
|
|
| 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: |
| 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()) |
|
|
