diff --git a/docs/pfi/model-quantization-playbook.md b/docs/pfi/model-quantization-playbook.md index 7b16943..7b1a3a5 100644 --- a/docs/pfi/model-quantization-playbook.md +++ b/docs/pfi/model-quantization-playbook.md @@ -226,6 +226,54 @@ bug — it is architectural). The proper upstream fix (vllm#51113) is in `main` Two cross-frontier peers (dvalin/bil-smithy) confirmed the bug class and pointed at the open symptom-twin issue #47087. +### 3.15 ⭐⭐ Fused 3-D MoE experts are INVISIBLE to a `targets=["Linear"]` recipe + +**Symptom: none.** The quant completes, the artifact loads, and 88.5% of the +model is still BF16. Nothing warns you. + +Modern MoE checkpoints store each layer's experts as **two fused 3-D +`nn.Parameter` tensors**, not as N `nn.Linear` modules. Gemma-4 26B-A4B: + + model.language_model.layers.N.experts.gate_up_proj BF16 [128, 1408, 2816] + model.language_model.layers.N.experts.down_proj BF16 [128, 2816, 704] + +Note the **absent `.weight` suffix** — that is the tell. `mlp.down_proj.weight` +is an `nn.Linear`; `experts.down_proj` is a bare parameter. + +Measured on that checkpoint, recipe targeting `["Linear"]`: + + Linear modules 427 + WILL quantize 205 (experts: 0) <- 22.84 B params untouched + +**This is the same defect that killed QLoRA on this architecture** — +`bitsandbytes` 4-bit replacement also walks `nn.Linear` modules and also +silently skipped the experts. Two different tools, one blind spot, because the +blind spot is in the *checkpoint layout*, not the tool. + +**The fix** (llm-compressor ≥ 0.12): + +```python +from llmcompressor.modeling.moe.linearize import linearize_moe +model = SomeForConditionalGeneration.from_pretrained(...) +linearize_moe(model) # BEFORE building the recipe +``` + + Linear modules 11947 + WILL quantize 11725 (experts: 11520) # 30 layers x 128 x 3 proj + +`linearize_moe` unfuses the 3-D parameters into per-expert +`experts.N.{gate,up,down}_proj` Linears. **No registration is needed** if the +module satisfies `FusedExpertsProtocol` structurally — bare `down_proj` plus +`gate_up_proj`/`up_proj` Parameters. `load_quantizable_moe(model_cls)` is the +faster variant that linearizes during load rather than after. + +**Always assert the expert count before spending GPU time** (§4.1). The +arithmetic is `layers × experts × projections`; if your target list does not +hit it exactly, the recipe is wrong and the failure is silent. + +⚠ **Keep routers in `ignore`.** A 4-bit router picks *different experts* — that +error does not average out downstream, it changes which weights run at all. + ### 3.4 Toolchain version deadlocks Both directions have burned us, so the resolution is: **use llm-compressor / compressed-tensors, diff --git a/scripts/erp-tune-serve/README.md b/scripts/erp-tune-serve/README.md new file mode 100644 index 0000000..8ec8439 --- /dev/null +++ b/scripts/erp-tune-serve/README.md @@ -0,0 +1,80 @@ +# ERP/RP tune → served NVFP4 seat + +Pipeline for turning the Gemma-4 26B-A4B ERP/RP LoRA into a servable NVFP4A16 +model on `ana-ml2`. Written 2026-08-24 alongside round 2 of the tune. + +Live copies run from `/tank/erp-tune/serve/` on ana-ml2. Model-agnostic quant +lessons belong in +[`docs/pfi/model-quantization-playbook.md`](../../docs/pfi/model-quantization-playbook.md); +the Gemma-4-specific ones are in +[`docs/pfi/gemma4-erp-tune-sizing.md`](../../docs/pfi/gemma4-erp-tune-sizing.md). + +## Order + +```bash +Q=/tank/aimodels/quant-work/.venv/bin/python # llmcompressor 0.12, ct 0.17.1 + +# 1. merge the adapter into bf16 (CPU, ~48 GB RAM, no GPU) +$Q merge_lora.py \ + --base /tank/aimodels/gemma4-26b-a4b-it-heretic-bf16 \ + --adapter /tank/erp-tune/run-01/adapter \ + --out /tank/erp-tune/serve/merged-bf16 + +# 2. PROVE the target set before spending GPU time +$Q quant_nvfp4a16.py --model /tank/erp-tune/serve/merged-bf16 \ + --out /tmp/x --calib-cache .jsonl --dry-run + +# 3. quantize +$Q quant_nvfp4a16.py --model /tank/erp-tune/serve/merged-bf16 \ + --out /tank/erp-tune/serve/nvfp4a16 \ + --calib-cache /tank/erp-tune/run-01/encode-cache/encoded-*.jsonl +``` + +## The three things that would silently ruin this + +**1. `targets=["Linear"]` misses every MoE expert.** Gemma-4 stores 128 experts +per layer as two fused 3-D `nn.Parameter`s, so a Linear-targeting recipe hits +205 of 427 modules and **zero** experts — 22.84 B params stay BF16 and nothing +warns you. `linearize_moe(model)` unfuses them (427 → 11,947 Linears, 11,520 +expert targets). Same blind spot that killed QLoRA here via `bitsandbytes`. +Playbook §3.15. **The dry run exists to catch this; use it.** + +**2. Shipping the base's own chat template is train/serve skew.** The trainee +base carries a *stale* 365-line `chat_template.jinja`; upstream's is 390. The +harness trained through upstream (config key `chat_template_path`), so the +merged model must ship upstream's. Wrong template presents as a tuning failure +with no error. `merge_lora.py` copies it explicitly and refuses if absent. + +**3. Calibration bakes a truncation cap into the tokenizer.** Playbook §3.14 — +a fast tokenizer called with `truncation=True` mutates its Rust backend state +in place, and `save_pretrained` persists it, clamping every prompt forever. +Sidestepped here by calibrating on the run's **encode cache** (already-tokenized +records) so the tokenizer is never called with truncation at all. Both scripts +still assert `tokenizer.json` has no `truncation` block before declaring success. + +## Why NVFP4**A16** and not the playbook's default mixed W4A4 + +Playbook §1 prefers mixed NVFP4-W4A4 + FP8. This seat deviates deliberately: + +- brokkr-smithy-dev benched the W4A4 quant of this checkpoint at **12% on + contradiction detection with CoT off against gen's 81%**, while T1/T3/T4/T5 + sat at 100%. Not general degradation — the signature of 4-bit *input + activations* on a reasoning-dense task. +- W4A4 KLD is 2–4× worse past ~10k ctx on sm_120; activation-quant noise + compounds with KV lookups. +- This is a 16,384-ctx RP seat. Long sessions **are** the workload. + +Cost accepted: A16 forces the Marlin kernel, ~half the prefill of native FP4. +Decode is memory-bound and barely moves. + +⚠ Several HF repos named `…-NVFP4A16` declare `input_activations num_bits 4` — +they are W4A4 wearing an A16 label. `quant_nvfp4a16.py` refuses if the emitted +config says `num_bits: 4`. Verify before substituting any upstream artifact. + +## Merge, don't hot-swap + +LoRA-on-NVFP4 hot-swap was a silent no-op on vLLM 0.24.0 (#47639, proven +quant-agnostic). Merging first means the quantizer sees ordinary bf16 weights +and the served artifact needs no adapter machinery. `merge_lora.py` asserts the +merge actually changed sampled target weights — a bit-identical merge would +otherwise ship the base model wearing the tune's name. diff --git a/scripts/erp-tune-serve/merge_lora.py b/scripts/erp-tune-serve/merge_lora.py new file mode 100755 index 0000000..7c77a99 --- /dev/null +++ b/scripts/erp-tune-serve/merge_lora.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 +"""Merge the ERP LoRA adapter into the bf16 base, producing servable weights. + +WHY MERGE RATHER THAN HOT-SWAP. Serving NVFP4 base + LoRA at runtime was a +silent no-op on vLLM 0.24.0 (#47639, proven quant-agnostic). Merging first +sidesteps it entirely: the quantizer then sees ordinary bf16 weights and the +served artifact needs no adapter machinery at all. + +⚠⚠ CHAT TEMPLATE. The trainee base ships a STALE 365-line chat_template.jinja; +upstream's is 390 lines. The harness deliberately trained through the UPSTREAM +template (config key `chat_template_path`), so the merged model MUST ship that +same upstream template. Shipping the base's own template here would be +train/serve skew with no error — it presents as a tuning failure. + +⚠ CPU merge. device_map=None keeps the 48 GiB on host RAM (566 GB total here) +so this can run while GPU0 is training. Do not use device_map="auto". + +⚠ Loader class. This checkpoint is Gemma4ForConditionalGeneration (vision + +audio towers present). Loading it as a plain CausalLM is playbook §3.2 — a +silent weight-load failure. +""" +import argparse +import json +import shutil +import sys +from pathlib import Path + +UPSTREAM_TEMPLATE = "/tank/aimodels/gemma4-26b-a4b-it-bf16/chat_template.jinja" + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--base", required=True) + ap.add_argument("--adapter", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--chat-template", default=UPSTREAM_TEMPLATE) + 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 peft import PeftModel + + print(f"[merge] loading base on CPU: {a.base}", flush=True) + model = Gemma4ForConditionalGeneration.from_pretrained( + a.base, dtype=torch.bfloat16, device_map=None, trust_remote_code=True, + ) + + # Count LoRA-target params before/after as a merge-actually-happened check. + print(f"[merge] applying adapter: {a.adapter}", flush=True) + before = {n: p.detach().clone() for n, p in model.named_parameters() + if n.endswith("self_attn.q_proj.weight") + and ".language_model.layers.0." in n} + + model = PeftModel.from_pretrained(model, a.adapter, is_trainable=False) + n_lora = sum(1 for n, _ in model.named_parameters() if "lora_" in n) + print(f"[merge] adapter tensors seen: {n_lora}", flush=True) + if n_lora == 0: + print("REFUSING: adapter contributed 0 tensors", file=sys.stderr) + return 2 + + model = model.merge_and_unload() + print("[merge] merged", flush=True) + + # ⚠ Prove the merge changed weights. A no-op merge is the failure mode that + # ships a base model wearing the tune's name, and nothing else would catch it. + changed = 0 + for n, p in model.named_parameters(): + if n in before: + if not torch.equal(p.detach(), before[n]): + changed += 1 + if changed == 0: + print("REFUSING: merge produced BIT-IDENTICAL weights on sampled " + "LoRA-target modules — the adapter was inert or did not apply", + file=sys.stderr) + return 3 + print(f"[merge] verified {changed}/{len(before)} sampled target(s) changed", flush=True) + + out.mkdir(parents=True, exist_ok=True) + print(f"[merge] saving to {out}", flush=True) + model.save_pretrained(out, safe_serialization=True) + + # Tokenizer straight from the base — never one that has been through + # calibration (playbook §3.14). + AutoTokenizer.from_pretrained(a.base, trust_remote_code=True).save_pretrained(out) + + # ⚠ Ship the UPSTREAM chat template, matching what training rendered. + src = Path(a.chat_template) + if not src.exists(): + print(f"REFUSING: chat template missing at {src}", file=sys.stderr) + return 4 + shutil.copy2(src, out / "chat_template.jinja") + n_lines = len(src.read_text().splitlines()) + print(f"[merge] chat_template.jinja <- {src} ({n_lines} lines)", 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", file=sys.stderr) + return 5 + + print(f"[merge] DONE -> {out}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/erp-tune-serve/quant_nvfp4a16.py b/scripts/erp-tune-serve/quant_nvfp4a16.py new file mode 100755 index 0000000..ff56014 --- /dev/null +++ b/scripts/erp-tune-serve/quant_nvfp4a16.py @@ -0,0 +1,190 @@ +#!/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) + oneshot( + model=model, dataset=ds, recipe=recipe, + 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) + src_tpl = Path(a.model) / "chat_template.jinja" + if src_tpl.exists(): + (out / "chat_template.jinja").write_text(src_tpl.read_text()) + print("[post] chat_template.jinja carried over", 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())