diff --git a/services/heretic2-nvfp4-quant/README.md b/services/heretic2-nvfp4-quant/README.md new file mode 100644 index 0000000..d521f96 --- /dev/null +++ b/services/heretic2-nvfp4-quant/README.md @@ -0,0 +1,48 @@ +# heretic2-nvfp4-quant — fast char-rp-reasoning seat (NVFP4 + MTP) + +Local NVFP4 quant of **NEO-CODE = Heretic2-Thinking** (Qwen3.6-27B) with the +Qwen3.6 **MTP head grafted back**, for a ~2.5–4× faster vLLM/MTP +`char-rp-reasoning` seat (buys reasoning-budget headroom → better GM planning +inside soong's latency window). R36 fast-seat spike, 2026-07-14. + +Fleet-first **local** NVFP4 quant — every other fleet NVFP4 model is *pulled* +pre-quantized; Heretic2 has none published, so we quantize it. + +## Fire sequence +1. **`graft_mtp.py`** — graft the 15 base-Qwen3.6 MTP tensors into Heretic2 BF16. + CPU-only, no GPU window. (Heretic2's finetune dropped the head; config declares + `mtp_num_hidden_layers=1` but ships 0 `mtp.*` tensors — verified.) +2. **`quant_nvfp4.py`** — llm-compressor NVFP4, Linear only; **GDN/vision/lm-head/ + norms/MTP kept BF16** (robbatt deckard-nvfp4 recipe + MTP). NEEDS a freed + Blackwell GPU (~55 GB) + an llmcompressor env (run in a vLLM container: + `pip install llmcompressor` on `vllm/vllm-openai:v0.24.0`). + - baseline: `--calib-mode text --calib neuralmagic/calibration` (AEON control) + - production: `--calib-mode chat --calib <512-row mix>` (brokkr/Dvalin) — rows + rendered via `apply_chat_template(enable_thinking=True)` so the forward-pass + sees the qwen3_coder tool-call XML = the seat's native activations. +3. **serve** (pantheon-27b-mtp-nvfp4 pattern): `vllm --quantization compressed-tensors + --speculative-config '{"method":"qwen3_5_mtp","num_speculative_tokens":3}' + --reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice`. +4. **P00 acceptance** (brokkr): soong 9-tool k5 rig on the quant — must hold + ~0.967 / perfect `attach_tool`. This is the **authoritative #355 check, NOT KL** + (KL can pass while the structured-tool path regresses). + +## Gates +- **GPU window** — both ana-ml2 Blackwell GPUs run ~full; NVFP4 is Blackwell-only + (irv-ml1's Ampere can't). The ~30–60 min quant needs a brief off-peak window + freeing a GPU. *Operator's call.* +- **Production calib** — brokkr/Dvalin assembling the 512-row mix; the tool-call-XML + slice (128 rows, 53 `attach_tool`) is ready. The AEON-baseline is fireable now. + +## Artifacts (on ana-ml2) +- Heretic2 BF16 (target): `/tank/aimodels/huggingface/hub/models--DavidAU--Qwen3.6-27B-Heretic2-Uncensored-Finetune-Thinking` +- base Qwen3.6-27B (MTP source): `/tank/aimodels/huggingface/hub/models--Qwen--Qwen3.6-27B` (15 mtp.* tensors, shards 13+15) +- AEON-baseline calib: `neuralmagic/calibration` (HF) +- `soong-tools-v0.3.13.json` — live 9-tool schema (structure source-of-truth; + **calib uses the PREFIXED `bifrost.soong-lab.*` runtime names** the seat emits) +- `extract_soong_tools.py` — how that schema was pulled from the deployed backend + +## Serve target +Replaces the current llama.cpp GGUF `char-rp-reasoning` seat (~59.5 tok/s) once +P00 passes. Deckard stays staged as rollback; the GGUF seat is the fallback until +the NVFP4 seat is validated + cut over. diff --git a/services/heretic2-nvfp4-quant/graft_mtp.py b/services/heretic2-nvfp4-quant/graft_mtp.py new file mode 100644 index 0000000..faa7c92 --- /dev/null +++ b/services/heretic2-nvfp4-quant/graft_mtp.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +"""Graft the Qwen3.6-27B MTP (multi-token-prediction) head into Heretic2-Thinking. + +DavidAU's Heretic2-Uncensored-Finetune-Thinking finetune dropped the MTP head +(config declares mtp_num_hidden_layers=1 but ships ZERO mtp.* weight tensors), +so the base Qwen/Qwen3.6-27B's 15 MTP tensors must be grafted back in BF16 +before NVFP4 quantization — this is what lets the served seat use `qwen3_5_mtp` +speculative decoding (n=3), the pantheon-27b-mtp-nvfp4 pattern proven on our +Blackwell. (R36 fast-seat spike, brokkr calib spec 2026-07-14.) + +CPU-only safetensors surgery — no GPU / no quant-window needed. Space-efficient: +symlinks Heretic2's large shards, adds one small model-mtp.safetensors, and +writes a merged index. Idempotent-ish: refuses to clobber a non-empty OUT_DIR. + +Usage (on ana-ml2): + python3 graft_mtp.py \ + --heretic2 /tank/aimodels/huggingface/hub/models--DavidAU--Qwen3.6-27B-Heretic2-Uncensored-Finetune-Thinking/snapshots/ \ + --base /tank/aimodels/huggingface/hub/models--Qwen--Qwen3.6-27B/snapshots/ \ + --out /tank/aimodels/heretic2-mtp-bf16 +""" +import argparse +import json +import os +import shutil +import sys + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +MTP_SHARD = "model-mtp.safetensors" +# MTP-related config keys we ensure are present on the grafted model (copied from +# base if Heretic2's config is missing any). +MTP_CFG_KEYS = ("mtp_num_hidden_layers", "mtp_use_dedicated_embeddings") + + +def _is_mtp(name: str) -> bool: + n = name.lower() + return n.startswith("mtp.") or "mtp." in n or "nextn" in n + + +def _tensor_nbytes(t: torch.Tensor) -> int: + return t.numel() * t.element_size() + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--heretic2", required=True, help="Heretic2 BF16 snapshot dir (quant target)") + ap.add_argument("--base", required=True, help="Qwen/Qwen3.6-27B snapshot dir (MTP source)") + ap.add_argument("--out", required=True, help="output dir for the grafted BF16 model") + args = ap.parse_args() + + her, base, out = args.heretic2, args.base, args.out + for d in (her, base): + if not os.path.isfile(os.path.join(d, "model.safetensors.index.json")): + print(f"ERROR: no index.json in {d}", file=sys.stderr) + return 2 + if os.path.isdir(out) and os.listdir(out): + print(f"ERROR: {out} exists and is non-empty — refusing to clobber", file=sys.stderr) + return 2 + os.makedirs(out, exist_ok=True) + + her_idx = json.load(open(os.path.join(her, "model.safetensors.index.json"))) + base_idx = json.load(open(os.path.join(base, "model.safetensors.index.json"))) + + # sanity: Heretic2 must NOT already have MTP weights; base MUST. + her_mtp = [k for k in her_idx["weight_map"] if _is_mtp(k)] + base_mtp = [k for k in base_idx["weight_map"] if _is_mtp(k)] + if her_mtp: + print(f"ERROR: Heretic2 already has {len(her_mtp)} mtp tensors — graft not needed", file=sys.stderr) + return 2 + if not base_mtp: + print("ERROR: base has no mtp tensors — wrong source model", file=sys.stderr) + return 2 + print(f"grafting {len(base_mtp)} MTP tensors from base into Heretic2 ({len(her_idx['weight_map'])} tensors)") + + # 1) stage Heretic2: symlink big shards, copy everything else. + for fn in sorted(os.listdir(her)): + src = os.path.join(her, fn) + if not os.path.isfile(src): + continue + dst = os.path.join(out, fn) + if fn.endswith(".safetensors"): + os.symlink(os.path.realpath(src), dst) + elif fn != "model.safetensors.index.json": # index rewritten below + shutil.copy2(src, dst) + + # 2) pull the MTP tensors out of base's shards into one new shard (BF16 preserved). + base_mtp_shards = sorted({base_idx["weight_map"][k] for k in base_mtp}) + mtp_tensors = {} + for shard in base_mtp_shards: + with safe_open(os.path.join(base, shard), framework="pt") as f: + for name in f.keys(): + if _is_mtp(name): + mtp_tensors[name] = f.get_tensor(name) + assert len(mtp_tensors) == len(base_mtp), f"expected {len(base_mtp)} mtp tensors, got {len(mtp_tensors)}" + dtypes = {str(t.dtype) for t in mtp_tensors.values()} + print(f" extracted {len(mtp_tensors)} mtp tensors, dtypes={dtypes}") + save_file(mtp_tensors, os.path.join(out, MTP_SHARD), metadata={"format": "pt"}) + + # 3) merged index: Heretic2 weight_map + the mtp tensors -> the new shard. + new_map = dict(her_idx["weight_map"]) + added_bytes = 0 + for name, t in mtp_tensors.items(): + new_map[name] = MTP_SHARD + added_bytes += _tensor_nbytes(t) + meta = dict(her_idx.get("metadata", {})) + if "total_size" in meta: + meta["total_size"] = int(meta["total_size"]) + added_bytes + out_idx = {"metadata": meta, "weight_map": new_map} + json.dump(out_idx, open(os.path.join(out, "model.safetensors.index.json"), "w"), indent=2) + + # 4) ensure config has complete MTP config (copy from base if Heretic2 lacks a key). + cfg_path = os.path.join(out, "config.json") + cfg = json.load(open(cfg_path)) + base_cfg = json.load(open(os.path.join(base, "config.json"))) + + def _get(d, k): + return d.get(k, d.get("text_config", {}).get(k)) + + changed = [] + for k in MTP_CFG_KEYS: + if _get(cfg, k) is None and _get(base_cfg, k) is not None: + cfg[k] = _get(base_cfg, k) + changed.append(k) + if changed: + json.dump(cfg, open(cfg_path, "w"), indent=2) + print(f" config MTP keys filled from base: {changed or 'none (already complete)'}") + + total = len(new_map) + print(f"DONE -> {out}") + print(f" total tensors: {total} (Heretic2 {len(her_idx['weight_map'])} + MTP {len(mtp_tensors)})") + print(f" new shard: {MTP_SHARD} (+{added_bytes/1e9:.2f} GB)") + print(" next: NVFP4 quant (quant_nvfp4.py) — keeps mtp.*/GDN/vision/lm-head/norms in BF16") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/services/heretic2-nvfp4-quant/quant_nvfp4.py b/services/heretic2-nvfp4-quant/quant_nvfp4.py new file mode 100644 index 0000000..6ae012d --- /dev/null +++ b/services/heretic2-nvfp4-quant/quant_nvfp4.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""NVFP4 quantize the MTP-grafted Heretic2 model (llm-compressor / compressed-tensors). + +Runs AFTER graft_mtp.py. Uses the fleet-proven compressed-tensors NVFP4 path +(the pantheon-27b-mtp-nvfp4 serving pattern already works on our Blackwell) with +the ignore-list from robbatt's on-fleet deckard-nvfp4 recipe PLUS the MTP head: +keep GDN/linear-attn, vision tower, lm-head, all norms, AND mtp.* in BF16; +NVFP4 only the dense Linear layers. (R36 fast-seat spike, brokkr calib spec.) + +⚠️ NEEDS: (a) an env with llmcompressor + a CUDA torch (run inside a vLLM +container: `pip install llmcompressor` on vllm/vllm-openai:v0.24.0), (b) a freed +Blackwell GPU (~55 GB — both ana-ml2 GPUs are normally full; needs an off-peak +window). API validated against llm-compressor at run time — treat the exact +symbol names as first-draft until a dry import confirms them. + +Two calib modes (brokkr's two artifacts): + --calib-mode text : AEON-baseline control (neuralmagic/calibration LLM split) + --calib-mode chat : production 512-row mix (JSONL rows {messages, tools}); + each row rendered via apply_chat_template(enable_thinking=True) + so the forward-pass sees the qwen3_coder tool-call XML = + the seat's native activations (the #355-preservation point). + +Usage: + python3 quant_nvfp4.py --model /tank/aimodels/heretic2-mtp-bf16 \ + --calib-mode text --calib neuralmagic/calibration --num-samples 160 \ + --out /tank/aimodels/heretic2-mtp-nvfp4-baseline + python3 quant_nvfp4.py --model /tank/aimodels/heretic2-mtp-bf16 \ + --calib-mode chat --calib /path/to/production_calib_512.jsonl \ + --out /tank/aimodels/heretic2-mtp-nvfp4-prod +""" +import argparse +import json +import sys + +# Ignore list = robbatt deckard-nvfp4 recipe + MTP head (brokkr: keep mtp.* BF16). +# Keeps in BF16: lm-head, embeddings, vision tower, GDN/linear-attn, all norms, MTP. +IGNORE = [ + "lm_head", + "re:.*embed_tokens$", + "re:visual.*", + "re:model.visual.*", + "re:.*linear_attn.*", + "re:.*norm.*", + "re:.*q_norm.*", + "re:.*k_norm.*", + "re:mtp.*", # MTP head stays BF16 for the qwen3_5_mtp spec-decode head +] + + +def load_calib_text(name, tokenizer, n, seqlen): + from datasets import load_dataset + ds = load_dataset(name, split="train").shuffle(seed=42).select(range(n)) + col = "text" if "text" in ds.column_names else ds.column_names[0] + return [tokenizer(x[col], truncation=True, max_length=seqlen) for x in ds] + + +def load_calib_chat(path, tokenizer, seqlen): + """Render {messages, tools} JSONL rows through the chat template with thinking on. + The assistant tool_calls (OpenAI form) serialize to qwen3_coder XML here, so the + calibration forward-pass sees the EXACT tool-call token distribution the seat emits.""" + rows = [json.loads(l) for l in open(path) if l.strip()] + out = [] + for r in rows: + text = tokenizer.apply_chat_template( + r["messages"], tools=r.get("tools"), + tokenize=False, add_generation_prompt=False, enable_thinking=True, + ) + out.append(tokenizer(text, truncation=True, max_length=seqlen)) + return out + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True, help="grafted BF16 model dir (graft_mtp.py output)") + ap.add_argument("--calib-mode", choices=["text", "chat"], required=True) + ap.add_argument("--calib", required=True, help="HF dataset name (text) or JSONL path (chat)") + ap.add_argument("--out", required=True) + ap.add_argument("--num-samples", type=int, default=512) + ap.add_argument("--seqlen", type=int, default=8192) + args = ap.parse_args() + + from transformers import AutoModelForCausalLM, AutoTokenizer + from llmcompressor import oneshot + from llmcompressor.modifiers.quantization import QuantizationModifier + + print(f"loading grafted model: {args.model}") + model = AutoModelForCausalLM.from_pretrained( + args.model, torch_dtype="auto", device_map="auto", trust_remote_code=True, + ) + tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + + print(f"building calibration ({args.calib_mode}, {args.num_samples} samples, seq {args.seqlen})") + if args.calib_mode == "text": + calib = load_calib_text(args.calib, tok, args.num_samples, args.seqlen) + else: + calib = load_calib_chat(args.calib, tok, args.seqlen) + print(f" {len(calib)} calibration rows") + + recipe = QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=IGNORE) + + print("running NVFP4 oneshot (Linear-only; GDN/vision/lm-head/norms/MTP kept BF16)") + oneshot( + model=model, dataset=calib, recipe=recipe, + max_seq_length=args.seqlen, num_calibration_samples=len(calib), + ) + + print(f"saving -> {args.out}") + model.save_pretrained(args.out, save_compressed=True) + tok.save_pretrained(args.out) + print("DONE. serve: vllm --quantization compressed-tensors " + "--speculative-config '{\"method\":\"qwen3_5_mtp\",\"num_speculative_tokens\":3}' " + "--reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice") + print("NEXT: ping brokkr -> P00 rig (soong 9-tool k5); acceptance = hold ~0.967/perfect attach_tool") + return 0 + + +if __name__ == "__main__": + sys.exit(main())