esh-pfi-infrastructure/services/erp-seat-quant/quant_nvfp4a16_gemma4_moe.py

#!/usr/bin/env python3
"""NVFP4A16 (weight-only) quant of a Gemma-4 26B-A4B **MoE** checkpoint for vLLM (compressed-tensors).

Built for the ERP-seat tunes (merged LoRA on Gemma-4-26B-A4B-it or its abliteration). Replicates
the published `prithivMLmods/gemma-4-26B-A4B-it-NVFP4A16` recipe (targets=Linear, NVFP4A16,
routers + vision + lm_head ignored) with the fleet's own calibration corpus and chat template.

Playbook rules honoured (docs/pfi/model-quantization-playbook.md):
  §3.15  fused 3-D MoE experts are INVISIBLE to targets=["Linear"] -> linearize_moe() FIRST,
         then ASSERT the expert Linear count (layers x experts x 3) before any GPU time.
  §3.15  routers stay BF16 (a 4-bit router picks different experts).
  §3.5   vision/audio towers + projector ignored (BF16); processor configs restored after save.
  §3.6   CPU-resident load (device_map=None); llm-compressor onloads one layer at a time.
  §3.10  do NOT set PYTORCH_CUDA_ALLOC_CONF=expandable_segments (corrupts retained tensors).
  §3.14  calibration bakes a truncation cap into tokenizer.json -> reset to null after save.
  §1     W4A16 chosen over mixed W4A4: RP long-session fidelity > prefill speed (gate-judged seat).
"""
import argparse, json, os, re, shutil, sys, hashlib

IGNORE = [
    "lm_head",
    "re:.*embed_tokens.*",
    "re:.*embed_vision.*",
    "re:.*vision_tower.*",
    "re:.*audio_tower.*",
    "re:.*audio.*",
    "re:.*multi_modal_projector.*",
    "re:.*mm_projector.*",
    "re:.*patch_embedder.*",
    "re:.*norm.*",
    "re:.*router.*",          # MoE routers stay BF16 (§3.15)
    "re:.*layer_scalar.*",
]

def _ignored(name):
    for pat in IGNORE:
        if pat.startswith("re:"):
            if re.fullmatch(pat[3:], name): return True
        elif name == pat or name.endswith("." + pat): return True
    return False

def enumerate_targets(model):
    import torch.nn as nn
    lin = [n for n, m in model.named_modules() if isinstance(m, nn.Linear)]
    will = [n for n in lin if not _ignored(n)]
    experts = [n for n in will if ".experts." in n]
    skipped = [n for n in lin if _ignored(n)]
    return lin, will, experts, skipped

def build_calib(path, tok, seqlen, n):
    from datasets import Dataset
    rows = [json.loads(l) for l in open(path) if l.strip()][:n]
    out = []
    for r in rows:
        msgs = []
        for m in r.get("messages", []):
            c = m.get("content")
            if isinstance(c, list):
                c = " ".join(p.get("text", "") for p in c if isinstance(p, dict))
            if c: msgs.append({"role": m.get("role", "user"), "content": c})
        if not msgs: continue
        try:
            text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
        except Exception:
            text = "\n".join(f"{m['role']}: {m['content']}" for m in msgs)
        out.append(tok(text, truncation=True, max_length=seqlen))
    return Dataset.from_list(out)

def sha256(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""): h.update(chunk)
    return h.hexdigest()

def post_steps(src, out):
    """§4.3-style post-steps for a Gemma-4 (no MTP head): processor configs, template, tokenizer cap, ignore check."""
    for f in ("processor_config.json", "preprocessor_config.json", "video_preprocessor_config.json", "generation_config.json"):
        s = os.path.join(src, f)
        if os.path.exists(s) and not os.path.exists(os.path.join(out, f)):
            shutil.copy(s, os.path.join(out, f)); print(f"[post] restored {f}")
    pc = os.path.join(out, "processor_config.json"); pp = os.path.join(out, "preprocessor_config.json")
    if os.path.exists(pc) and not os.path.exists(pp):
        d = json.load(open(pc))
        if "image_processor" in d:
            json.dump(dict(d["image_processor"]), open(pp, "w"), indent=1); print("[post] materialized preprocessor_config.json from processor_config.image_processor")
    # chat template: ship the SOURCE's (the one training/serving used), byte-identical
    st = os.path.join(src, "chat_template.jinja"); ot = os.path.join(out, "chat_template.jinja")
    shutil.copy(st, ot); print(f"[post] chat_template.jinja <- source, sha256 {sha256(ot)[:16]}")
    # tokenizer truncation cap (§3.14)
    tj = os.path.join(out, "tokenizer.json"); t = json.load(open(tj))
    if t.get("truncation") is not None:
        print(f"[post] ⚠ tokenizer.json had truncation={t['truncation']} baked in -> reset to null")
        t["truncation"] = None; json.dump(t, open(tj, "w"), ensure_ascii=False)
    else:
        print("[post] tokenizer.json truncation: null (clean)")
    # quantization_config ignore must still carry the routers + vision
    cfg = json.load(open(os.path.join(out, "config.json"))); ig = cfg["quantization_config"].get("ignore", [])
    has_router = any("router" in x for x in ig); has_vision = any("vision" in x for x in ig)
    print(f"[post] quantization_config.ignore: {len(ig)} entries; routers={has_router} vision={has_vision}")
    if not (has_router and has_vision):
        print("[post] ⚠ REFUSING: ignore list lost routers or vision — llm-compressor pruned unmatched entries; investigate before serving", file=sys.stderr)
        return 3
    return 0

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--calib", default="/tank/aimodels/heretic2-nvfp4-work/production_calib_512.jsonl")
    ap.add_argument("--num-samples", type=int, default=256); ap.add_argument("--seqlen", type=int, default=8192)
    ap.add_argument("--scheme", default="NVFP4A16")
    ap.add_argument("--expect-experts", type=int, default=30 * 128 * 3, help="layers x experts x projections; assert before GPU time")
    ap.add_argument("--dry-run", action="store_true", help="load + linearize + enumerate targets only (no GPU, no save)")
    a = ap.parse_args()

    from transformers import AutoTokenizer, Gemma4ForConditionalGeneration
    print(f"[load] {a.model} (CPU-resident)", flush=True)
    tok = AutoTokenizer.from_pretrained(a.model)
    model = Gemma4ForConditionalGeneration.from_pretrained(a.model, torch_dtype="auto", device_map=None)
    from llmcompressor.modeling.moe.linearize import linearize_moe
    linearize_moe(model)
    lin, will, experts, skipped = enumerate_targets(model)
    print(f"[targets] Linear modules {len(lin)}  WILL quantize {len(will)}  (experts: {len(experts)})  ignored {len(skipped)}", flush=True)
    print("[targets] ignored sample:", sorted({re.sub(r'\.\d+\.', '.N.', n) for n in skipped})[:20], flush=True)
    print("[targets] quantized sample:", sorted({re.sub(r'\.\d+\.', '.N.', n) for n in will})[:12], flush=True)
    if len(experts) != a.expect_experts:
        print(f"[targets] ⚠ REFUSING: expert Linear count {len(experts)} != expected {a.expect_experts} (§3.15)", file=sys.stderr); return 2
    if any("router" in n or "vision" in n or "audio" in n for n in will):
        print("[targets] ⚠ REFUSING: a router/vision/audio Linear is in the quantize set", file=sys.stderr); return 2
    if a.dry_run:
        print("[dry-run] OK — targets proven; exiting before calibration"); return 0

    print(f"[calib] <= {a.num_samples} samples @ seq {a.seqlen} from {a.calib}", flush=True)
    calib = build_calib(a.calib, tok, a.seqlen, a.num_samples); print(f"[calib] {len(calib)} rows", flush=True)
    from llmcompressor import oneshot
    from llmcompressor.modifiers.quantization import QuantizationModifier
    recipe = QuantizationModifier(targets="Linear", scheme=a.scheme, ignore=IGNORE)
    print(f"[quant] oneshot scheme={a.scheme} targets=Linear (post-linearize)", flush=True)
    oneshot(model=model, dataset=calib, recipe=recipe, num_calibration_samples=len(calib), max_seq_length=a.seqlen)
    print(f"[save] -> {a.out}", flush=True)
    model.save_pretrained(a.out, save_compressed=True); tok.save_pretrained(a.out)
    rc = post_steps(a.model, a.out)
    print("DONE" if rc == 0 else f"DONE WITH POST-STEP FAILURE rc={rc}", flush=True); return rc

if __name__ == "__main__":
    sys.exit(main())