feat(heretic2-nvfp4): MTP-graft + NVFP4 quant scripts + pipeline README (fire-ready)
graft_mtp.py: grafts the 15 base-Qwen3.6 MTP tensors into Heretic2 BF16 (CPU-only). quant_nvfp4.py: llm-compressor NVFP4 (Linear only; GDN/vision/lm-head/norms/MTP kept BF16 per robbatt's deckard recipe + brokkr's spec); text (AEON-baseline) or chat (production, apply_chat_template renders qwen3_coder XML) calib modes. README: fire sequence + gates (GPU window, production calib) + artifacts. Spike gated only on: (1) off-peak Blackwell GPU window, (2) brokkr's production calib.
This commit is contained in:
@@ -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.
|
||||||
@@ -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/<hash> \
|
||||||
|
--base /tank/aimodels/huggingface/hub/models--Qwen--Qwen3.6-27B/snapshots/<hash> \
|
||||||
|
--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())
|
||||||
@@ -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())
|
||||||
Reference in New Issue
Block a user