fix: harden quantized checkpoint integration

scripts/smoke_load_quant.py

@@ -1,55 +1,100 @@
-"""Smoke test: load a quantized checkpoint through the OBLITERATUS loader.
+"""Smoke-load a quantized checkpoint through the OBLITERATUS loader."""
 
-Verifies: detection, dequantization to BF16, no quantized tensors left in
-the model, config stripped of quantization_config, and a short generation.
-"""
+from __future__ import annotations
 
-import sys
+import argparse
 
 import torch
 
 from obliteratus.models import quant_dequant as qd
 from obliteratus.models.loader import load_model
 
-repo = sys.argv[1]
 
-det = qd.detect_quant_scheme(repo)
-print(f"[detection] {det.scheme.value} (inverse={det.scale_is_inverse}, group={det.group_size})")
-assert det.scheme not in (qd.QuantScheme.NONE, qd.QuantScheme.UNSUPPORTED), det
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("repo", help="Local checkpoint path or Hugging Face repository")
+    parser.add_argument(
+        "--revision",
+        help="Immutable Hub commit, tag, or branch to load",
+    )
+    parser.add_argument(
+        "--local-files-only",
+        action="store_true",
+        help="Refuse network access and use only locally cached files",
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help=(
+            "Explicitly allow checkpoint-provided Python code to execute. "
+            "Off by default; review and pin the repository revision first."
+        ),
+    )
+    return parser
 
-handle = load_model(
-    repo,
-    task="causal_lm",
-    device="auto",
-    dtype="bfloat16",
-    trust_remote_code=True,
-)
-model = handle.model
-print(f"[load] ok — scheme tag: {getattr(model, '_obliteratus_dequantized_scheme', None)}")
 
-bad = []
-n_params = 0
-for name, p in model.named_parameters():
-    n_params += 1
-    if p.dtype in qd.FP8_DTYPES or p.dtype == torch.uint8:
-        bad.append((name, str(p.dtype)))
-    if torch.isnan(p.data).any().item() if p.data.is_floating_point() else False:
-        bad.append((name, "NaN"))
-print(f"[check] {n_params} params, quantized/NaN leftovers: {bad[:10] or 'NONE'}")
-assert not bad, bad
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    det = qd.detect_quant_scheme(
+        args.repo,
+        revision=args.revision,
+        local_files_only=args.local_files_only,
+    )
+    print(
+        f"[detection] {det.scheme.value} "
+        f"(block_inverse={det.scale_is_inverse}, "
+        f"global_inverse={det.global_scale_is_inverse}, group={det.group_size})"
+    )
+    if det.scheme in (qd.QuantScheme.NONE, qd.QuantScheme.UNSUPPORTED):
+        raise RuntimeError(det.reason or f"checkpoint is not FP8/NVFP4: {args.repo}")
 
-assert getattr(model.config, "quantization_config", None) is None, "quantization_config survived"
+    handle = load_model(
+        args.repo,
+        task="causal_lm",
+        device="auto",
+        dtype="bfloat16",
+        trust_remote_code=args.trust_remote_code,
+        revision=args.revision,
+        local_files_only=args.local_files_only,
+    )
+    model = handle.model
+    print(
+        "[load] ok — scheme tag: "
+        f"{getattr(model, '_obliteratus_dequantized_scheme', None)}"
+    )
 
-tok = handle.tokenizer
-prompt = "The capital of France is"
-try:
-    inputs = tok(prompt, return_tensors="pt").to(model.device)
-except Exception:
-    # Omni processors: fall back to bare tokenizer encode
-    ids = tok.encode(prompt, return_tensors="pt").to(model.device)
-    inputs = {"input_ids": ids}
-with torch.no_grad():
-    out = model.generate(**inputs, max_new_tokens=16, do_sample=False)
-text = tok.decode(out[0][-16:] if out.dim() > 1 else out[-16:], skip_special_tokens=True)
-print(f"[generate] {text!r}")
-print("SMOKE_LOAD_OK")
+    bad = []
+    n_params = 0
+    for name, param in model.named_parameters():
+        n_params += 1
+        if param.dtype in qd.FP8_DTYPES or param.dtype == torch.uint8:
+            bad.append((name, str(param.dtype)))
+        if param.data.is_floating_point() and torch.isnan(param.data).any().item():
+            bad.append((name, "NaN"))
+    print(f"[check] {n_params} params, quantized/NaN leftovers: {bad[:10] or 'NONE'}")
+    if bad:
+        raise RuntimeError(f"quantized or non-finite tensors remain: {bad[:10]}")
+
+    if getattr(model.config, "quantization_config", None) is not None:
+        raise RuntimeError("quantization_config survived dequantization")
+
+    tokenizer = handle.tokenizer
+    prompt = "The capital of France is"
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    except Exception:
+        ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+        inputs = {"input_ids": ids}
+    with torch.no_grad():
+        output = model.generate(**inputs, max_new_tokens=16, do_sample=False)
+    text = tokenizer.decode(
+        output[0][-16:] if output.dim() > 1 else output[-16:],
+        skip_special_tokens=True,
+    )
+    print(f"[generate] {text!r}")
+    print("SMOKE_LOAD_OK")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())