mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
"""Smoke-load a quantized checkpoint through the OBLITERATUS loader."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from obliteratus.models import quant_dequant as qd
|
|
from obliteratus.models.loader import load_model
|
|
|
|
|
|
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
|
|
|
|
|
|
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}")
|
|
|
|
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)}"
|
|
)
|
|
|
|
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())
|