File size: 5,300 Bytes
9589849 | 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 133 134 135 136 137 138 139 140 | #!/usr/bin/env python3
"""Validate an HF weight export and its terminal Prime-RL metrics."""
from __future__ import annotations
import argparse
import json
import math
from collections import defaultdict
from pathlib import Path
def merged_metrics(path: Path) -> dict[int, dict]:
by_step: dict[int, dict] = defaultdict(dict)
for line in path.read_text().splitlines():
if line.strip():
row = json.loads(line)
by_step[int(row["step"])].update(row)
return dict(by_step)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("model", type=Path)
parser.add_argument("metrics", type=Path)
parser.add_argument("--expected-step", type=int, required=True)
parser.add_argument(
"--scan-finite",
action="store_true",
help="load every floating tensor and reject NaN/Inf (slower)",
)
args = parser.parse_args()
model = args.model.resolve()
problems: list[str] = []
if not (model / "STABLE").is_file():
problems.append("missing STABLE marker")
index_path = model / "model.safetensors.index.json"
config_path = model / "config.json"
if not index_path.is_file():
problems.append("missing model.safetensors.index.json")
if not config_path.is_file():
problems.append("missing config.json")
if problems:
print(json.dumps({"healthy": False, "problems": problems}, indent=2))
raise SystemExit(1)
index = json.loads(index_path.read_text())
weight_map = index.get("weight_map") or {}
if not weight_map:
problems.append("empty weight_map")
expected_by_shard: dict[str, set[str]] = defaultdict(set)
for key, shard in weight_map.items():
expected_by_shard[str(shard)].add(str(key))
from safetensors import safe_open
tensor_count = 0
parameter_elements = 0
scanned_elements = 0
dtypes: dict[str, int] = defaultdict(int)
for shard_name, expected_keys in sorted(expected_by_shard.items()):
shard_path = model / shard_name
if not shard_path.is_file() or shard_path.stat().st_size == 0:
problems.append(f"missing or empty shard: {shard_name}")
continue
try:
with safe_open(shard_path, framework="pt", device="cpu") as handle:
actual_keys = set(handle.keys())
if actual_keys != expected_keys:
missing = len(expected_keys - actual_keys)
extra = len(actual_keys - expected_keys)
problems.append(
f"index/header mismatch in {shard_name}: missing={missing}, extra={extra}"
)
for key in sorted(actual_keys):
view = handle.get_slice(key)
shape = tuple(view.get_shape())
elements = math.prod(shape)
tensor_count += 1
parameter_elements += elements
dtypes[str(view.get_dtype())] += 1
if args.scan_finite:
tensor = handle.get_tensor(key)
if tensor.is_floating_point() or tensor.is_complex():
scanned_elements += tensor.numel()
if not tensor.isfinite().all().item():
problems.append(f"non-finite tensor: {key}")
except Exception as error: # a corrupt header/shard must fail the gate
problems.append(f"cannot load {shard_name}: {type(error).__name__}: {error}")
by_step = merged_metrics(args.metrics)
if args.expected_step not in by_step:
problems.append(f"metrics do not contain expected step {args.expected_step}")
terminal = {}
else:
terminal = by_step[args.expected_step]
for field in ("loss/mean", "optim/grad_norm", "optim/lr"):
value = terminal.get(field)
if not isinstance(value, (int, float)) or not math.isfinite(value):
problems.append(f"terminal metric {field} is not finite: {value!r}")
if terminal.get("loss/nan_count", 0) != 0:
problems.append(
f"terminal loss/nan_count is {terminal.get('loss/nan_count')!r}"
)
result = {
"healthy": not problems,
"model": str(model),
"expected_step": args.expected_step,
"stable_marker": (model / "STABLE").is_file(),
"shards": len(expected_by_shard),
"tensors": tensor_count,
"tensor_elements": parameter_elements,
"dtype_tensor_counts": dict(sorted(dtypes.items())),
"finite_scan_enabled": args.scan_finite,
"finite_scanned_elements": scanned_elements,
"terminal_metrics": {
key: terminal.get(key)
for key in (
"loss/mean",
"loss/nan_count",
"optim/grad_norm",
"optim/lr",
"perf/throughput",
"perf/peak_memory",
"progress/num_tokens",
"progress/num_samples",
)
},
"problems": problems,
}
print(json.dumps(result, indent=2))
if problems:
raise SystemExit(1)
if __name__ == "__main__":
main()
|