NVFP4 (ModelOpt) quantization of sensenova/SenseNova-Loop-9B-Preview @27abea77 with BF16 MTP head
a57aa9f verified | #!/usr/bin/env python3 | |
| """HF-transformers functional checks for SenseNova-Loop (reference BF16 backbone + FP32 mHC). | |
| Checks: mHC router dtype, shared-weight looping, text greedy, thinking, tool-call format, | |
| image understanding, and cached-vs-uncached decoding consistency (teacher-forced argmax). | |
| Writes a JSON report; does not modify the model directory. | |
| usage: python -I scripts/10_reference_checks.py MODEL_DIR OUT_JSON [IMAGE] | |
| """ | |
| import json, sys, time | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| model_dir, out_path = sys.argv[1], sys.argv[2] | |
| image_path = sys.argv[3] if len(sys.argv) > 3 else None | |
| dev = "cuda:0" | |
| report = {"model_dir": model_dir, "torch": torch.__version__} | |
| t0 = time.time() | |
| proc = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True) | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| model_dir, trust_remote_code=True, dtype=torch.bfloat16, device_map={"": dev}).eval() | |
| report["load_s"] = round(time.time() - t0, 1) | |
| report["gpu_mem_after_load_GiB"] = round(torch.cuda.memory_allocated() / 2**30, 2) | |
| text = model.model.language_model | |
| report["loop_plan_len"] = len(text._loop_plan) | |
| routers = {n: p.dtype for n, p in model.named_parameters() if ".mhc_" in n} | |
| report["mhc_param_count"] = len(routers) | |
| report["mhc_all_fp32"] = all(d == torch.float32 for d in routers.values()) | |
| report["layer_param_dtypes"] = sorted({str(p.dtype) for n, p in model.named_parameters() if ".mhc_" not in n}) | |
| # Each physical looped layer must be called twice per forward (shared weights). | |
| calls = {} | |
| hooks = [l.register_forward_hook(lambda m, i, o, k=k: calls.__setitem__(k, calls.get(k, 0) + 1)) | |
| for k, l in enumerate(text.layers)] | |
| tok = proc.tokenizer | |
| def chat(messages, max_new=128, tools=None, thinking=False, images=None): | |
| prompt = proc.apply_chat_template(messages, tools=tools, tokenize=False, add_generation_prompt=True, | |
| enable_thinking=thinking) | |
| inputs = proc(text=[prompt], images=images, return_tensors="pt").to(dev) | |
| t = time.time() | |
| with torch.inference_mode(): | |
| out = model.generate(**inputs, max_new_tokens=max_new, do_sample=False, use_cache=True) | |
| gen = out[0, inputs["input_ids"].shape[1]:] | |
| return {"prompt_tokens": int(inputs["input_ids"].shape[1]), "gen_tokens": int(gen.shape[0]), | |
| "seconds": round(time.time() - t, 2), "text": tok.decode(gen, skip_special_tokens=False)}, out, inputs | |
| r, _, _ = chat([{"role": "user", "content": "What is 13 times 7? Reply with the number only."}], 16) | |
| report["text_mult"] = r | |
| report["loop_calls_per_layer"] = {str(k): v for k, v in sorted(calls.items())} | |
| report["looped_layers_called_twice"] = all(calls[k] == (2 if 4 <= k < 28 else 1) * calls[0] for k in calls) | |
| for h in hooks: | |
| h.remove() | |
| report["text_haiku"] = chat([{"role": "user", "content": "Write a haiku about a tennis ball."}], 64)[0] | |
| report["thinking"] = chat([{"role": "user", "content": "If a train goes 60 km in 45 minutes, what is its speed in km/h?"}], | |
| 384, thinking=True)[0] | |
| tools = [{"type": "function", "function": { | |
| "name": "get_weather", "description": "Get current weather for a city.", | |
| "parameters": {"type": "object", "properties": { | |
| "city": {"type": "string"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}}, | |
| "required": ["city"]}}}] | |
| r = chat([{"role": "user", "content": "What's the weather in Paris in celsius?"}], 96, tools=tools)[0] | |
| r["has_tool_call_xml"] = ("<tool_call>" in r["text"] and "<function=get_weather>" in r["text"] | |
| and "<parameter=city>" in r["text"]) | |
| report["tool_call"] = r | |
| if image_path: | |
| img = Image.open(image_path).convert("RGB") | |
| msgs = [{"role": "user", "content": [{"type": "image"}, | |
| {"type": "text", "text": "Describe this image in two sentences."}]}] | |
| report["image"] = chat(msgs, 96, images=[img])[0] | |
| # Cached generation consistency: greedy tokens from cached decoding must equal the | |
| # argmax of a single uncached forward over prompt+generation (teacher forcing). | |
| r, out, inputs = chat([{"role": "user", "content": "Explain in one paragraph why the sky is blue."}], 64) | |
| plen = inputs["input_ids"].shape[1] | |
| with torch.inference_mode(): | |
| logits = model(input_ids=out, use_cache=False).logits[0, plen - 1:-1].float() | |
| pred = logits.argmax(-1) | |
| gen = out[0, plen:] | |
| report["cache_consistency"] = {"n": int(gen.numel()), "match": int((pred == gen).sum()), | |
| "first_mismatch": int((pred != gen).nonzero()[0]) if (pred != gen).any() else None, | |
| "finite_logits": bool(torch.isfinite(logits).all())} | |
| report["peak_gpu_mem_GiB"] = round(torch.cuda.max_memory_allocated() / 2**30, 2) | |
| json.dump(report, open(out_path, "w"), indent=2, ensure_ascii=False) | |
| print(json.dumps(report, indent=2, ensure_ascii=False)) | |