esh-pfi-infrastructure/scripts/erp-tune-serve/quant_nvfp4a16.py

#!/usr/bin/env python3
"""NVFP4A16 quantize the merged ERP/RP tune (Gemma-4 26B-A4B MoE).

SCHEME: weight-only NVFP4 **A16**, not W4A4. This is not the playbook's general
default (§1 prefers mixed NVFP4-W4A4 + FP8) and the deviation is deliberate and
measured on THIS architecture:

  * brokkr-smithy-dev benched the W4A4 serving quant of gemma4-26b-a4b-it and
    got **12% on contradiction detection with CoT off, against gen's 81%**,
    while T1/T3/T4/T5 all sat at 100%. Not general degradation - exactly the
    shape 4-bit INPUT ACTIVATIONS produce on the most reasoning-dense task.
  * NVIDIA moved to W4A16 for sm_120 long-context: W4A4 KLD is 2-4x worse past
    ~10k ctx, activation-quant noise compounding with KV lookups.
  * This seat is a 16,384-ctx RP model. Long sessions ARE the workload.

Cost paid on purpose: A16 forces the Marlin kernel, roughly half the prefill of
native FP4. Decode is memory-bound and barely moves. Accepted.

CALIBRATION uses the run's own encode cache - the exact token sequences the
model trained on, already rendered through the correct upstream chat template.
That is both maximally faithful AND sidesteps playbook 3.14 entirely, because
we never call the tokenizer with truncation=True at all.
"""
import argparse
import json
import sys
from pathlib import Path

# NVFP4 only the language-model dense Linears + the MoE experts.
# Everything here stays BF16.
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.*",
    # ⚠ Routers stay BF16. The shipped gemma4-26b-a4b-it-nvfp4 artifact ignores
    # every `router.proj`, and a 4-bit router picks different experts - the
    # error does not average out downstream, it changes which weights run.
    "re:.*router.*",
]


def build_calib(cache_path, n, seqlen):
    """Calibration set straight from the training encode cache.

    Records are already tokenized and already rendered through the upstream
    chat template, so this is the true training distribution. Long sequences
    matter more than sample count for long-context fidelity, so prefer the
    longest records rather than the first N.
    """
    from datasets import Dataset
    rows = []
    with open(cache_path) as fh:
        for line in fh:
            r = json.loads(line)
            rows.append(r["input_ids"])
    rows.sort(key=len, reverse=True)
    picked = rows[:n]
    out = [{"input_ids": ids[:seqlen],
            "attention_mask": [1] * len(ids[:seqlen])} for ids in picked]
    lens = [len(o["input_ids"]) for o in out]
    print("[calib] %d samples, tokens min/mean/max %d/%d/%d" % (
        len(out), min(lens), sum(lens) // len(lens), max(lens)), flush=True)
    return Dataset.from_list(out)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, help="merged bf16 model")
    ap.add_argument("--out", required=True)
    ap.add_argument("--calib-cache", required=True, help="encode-cache jsonl")
    ap.add_argument("--num-calib", type=int, default=256)
    ap.add_argument("--seqlen", type=int, default=16384)
    ap.add_argument("--dry-run", action="store_true",
                    help="resolve targets and print what WOULD be quantized, then exit")
    a = ap.parse_args()

    out = Path(a.out)
    if out.exists() and any(out.iterdir()):
        print(f"REFUSING: {out} exists and is non-empty", file=sys.stderr)
        return 1

    import torch
    from transformers import AutoTokenizer, Gemma4ForConditionalGeneration
    from llmcompressor import oneshot
    from llmcompressor.modifiers.quantization import QuantizationModifier
    from llmcompressor.modeling.moe.linearize import linearize_moe

    print(f"[load] {a.model} on CPU (oneshot onloads layer-by-layer)", flush=True)
    model = Gemma4ForConditionalGeneration.from_pretrained(
        a.model, dtype=torch.bfloat16, device_map=None, trust_remote_code=True)
    tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)

    # ⚠⚠ WITHOUT THIS THE MoE STAYS BF16. Gemma-4 stores each layer's 128
    # experts as two fused 3-D nn.Parameters (`gate_up_proj` [128,1408,2816],
    # `down_proj` [128,2816,704]) - NOT nn.Linear modules. A recipe targeting
    # ["Linear"] therefore matches 205 of 427 modules and ZERO experts, leaving
    # 22.84 B params (88.5% of the model) unquantized. That is precisely how
    # QLoRA failed on this architecture via bitsandbytes, reproduced in a
    # different tool.
    #
    # `linearize_moe` unfuses them into per-expert `experts.N.{gate,up,down}_proj`
    # Linear modules. Gemma-4 needs no registration - it satisfies
    # FusedExpertsProtocol structurally (bare `down_proj` + `gate_up_proj`
    # Parameters). Verified by the dry run: experts go 0 -> 11,520 targets.
    print("[moe] linearizing fused experts", flush=True)
    linearize_moe(model)

    # ⚠ §4.1 - prove the target set BEFORE spending GPU time. A recipe whose
    # ignore regexes silently miss the experts produces a "quantized" model
    # that is mostly still bf16, which is exactly how QLoRA failed on this
    # architecture (bitsandbytes skipped the fused 3-D expert params).
    import re as _re
    pats = [p[3:] for p in IGNORE if p.startswith("re:")]
    lits = [p for p in IGNORE if not p.startswith("re:")]
    lin = [n for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)]
    def ignored(n):
        return any(l in n for l in lits) or any(_re.search(p, n) for p in pats)
    tgt = [n for n in lin if not ignored(n)]
    exp = [n for n in tgt if ".experts." in n]
    rtr = [n for n in lin if "router" in n]
    print("[targets] Linear modules   %d" % len(lin))
    print("[targets] WILL quantize    %d  (experts: %d)" % (len(tgt), len(exp)))
    print("[targets] ignored          %d  (routers: %d)" % (len(lin) - len(tgt), len(rtr)))
    if exp == []:
        print("REFUSING: zero expert Linears targeted. The MoE would stay bf16 - "
              "this is the QLoRA failure mode. Check the model unfused its "
              "experts into experts.N.* modules.", file=sys.stderr)
        return 2
    for n in tgt[:3] + exp[:2]:
        print("   +", n)
    if a.dry_run:
        print("[dry-run] stopping before quantization")
        return 0

    ds = build_calib(a.calib_cache, a.num_calib, a.seqlen)

    recipe = QuantizationModifier(
        targets=["Linear"], scheme="NVFP4A16", ignore=IGNORE,
    )

    print("[oneshot] starting", flush=True)
    # ⚠ `processor` must be passed EXPLICITLY. This is a multimodal
    # (Gemma4ForConditionalGeneration) checkpoint, and llmcompressor's
    # auto-init fails on it with "An error occurred when attempting to
    # initialize model processor, which is required when a dataset is
    # provided." Calibration here is text-only - the records come from the
    # training encode cache - so the tokenizer is the correct processor.
    oneshot(
        model=model, dataset=ds, recipe=recipe, processor=tok,
        max_seq_length=a.seqlen, num_calibration_samples=len(ds),
        output_dir=str(out),
    )
    print("[oneshot] done", flush=True)

    # ⚠ playbook 3.14 - NEVER ship the calibration tokenizer. Re-read pristine.
    AutoTokenizer.from_pretrained(a.model, trust_remote_code=True).save_pretrained(out)
    import shutil as _sh
    for aux in ("chat_template.jinja", "processor_config.json",
                "preprocessor_config.json", "video_preprocessor_config.json",
                "special_tokens_map.json", "generation_config.json"):
        src_aux = Path(a.model) / aux
        if src_aux.exists():
            _sh.copy2(src_aux, out / aux)
            print(f"[post] carried {aux}", flush=True)

    tj = out / "tokenizer.json"
    if tj.exists() and json.loads(tj.read_text()).get("truncation"):
        print("REFUSING: shipped tokenizer carries a truncation cap "
              "(playbook 3.14) - the seat would clamp every prompt",
              file=sys.stderr)
        return 3
    print("[post] tokenizer truncation: clean", flush=True)

    cfg = json.loads((out / "config.json").read_text())
    qc = cfg.get("quantization_config", {})
    print("[verify] quant_method %s  format %s" % (
        qc.get("quant_method"), qc.get("format")))
    for g, v in (qc.get("config_groups") or {}).items():
        ia = v.get("input_activations")
        print("[verify] %s: w=%s a=%s" % (
            g, (v.get("weights") or {}).get("num_bits"),
            (ia or {}).get("num_bits") if ia else "null (A16)"))
        if ia and ia.get("num_bits") == 4:
            print("REFUSING: input_activations num_bits=4 - this is W4A4 "
                  "wearing an A16 label. See the header for why that is wrong "
                  "for this seat.", file=sys.stderr)
            return 4
    print(f"[quant] DONE -> {out}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())