vh/esh-pfi-infrastructure
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

feat(erp-tune): NVFP4A16 serving pipeline, and the MoE landmine it uncovered

Browse Source 
Merge + quantize path for turning the Gemma-4 26B-A4B ERP/RP LoRA into a
servable NVFP4A16 seat, plus a playbook entry for the defect found while
validating it.

The landmine (playbook §3.15): a `targets=["Linear"]` NVFP4 recipe silently
misses every MoE expert on this architecture. Gemma-4 stores each layer's 128
experts as two fused 3-D nn.Parameter tensors, not nn.Linear modules, so the
recipe resolves 205 of 427 modules and ZERO experts — 22.84 B params, 88.5% of
the model, left in BF16 with no warning. This is the same blind spot that
killed QLoRA here via bitsandbytes; the tool changed, the checkpoint layout did
not.

  before linearize_moe:  427 Linears,    205 targeted, experts 0
  after  linearize_moe:  11,947 Linears, 11,725 targeted, experts 11,520
                         (30 layers x 128 experts x 3 projections)

llmcompressor's linearize_moe unfuses them; no registration needed because
Gemma-4 satisfies FusedExpertsProtocol structurally. Caught by an §4.1 dry run
that asserts the expert count before any GPU spend, which is now the documented
requirement rather than an optional step.

Scheme is NVFP4A16, deviating from the playbook's mixed-W4A4 default on
measured grounds: brokkr-smithy-dev benched the W4A4 quant of this checkpoint
at 12% on contradiction detection with CoT off against gen's 81%, the signature
of 4-bit input activations on a reasoning-dense task, and W4A4 KLD degrades
2-4x past ~10k ctx on sm_120. This is a 16,384-ctx RP seat. Marlin's prefill
cost is accepted.

Two further silent-failure guards, both from prior hard-won lessons:
- the merged model ships the UPSTREAM chat template, not the trainee base's
  stale 365-line one, because training rendered through upstream and the
  mismatch would present as a tuning failure
- calibration reads the run's own encode cache rather than re-tokenizing, which
  sidesteps §3.14 (a fast tokenizer mutated by truncation=True and persisted by
  save_pretrained clamps every prompt forever)

Merge-then-quantize rather than LoRA hot-swap, since hot-swap onto NVFP4 was a
silent no-op on vLLM 0.24.0 (#47639). merge_lora.py asserts sampled target
weights actually changed, so an inert adapter cannot ship as a tune.
This commit is contained in:
vh
2026-08-24 23:11:33 -07:00
parent 7b5fd91d3c
commit 6a8582936e
4 changed files with 428 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/erp-tune-serve/README.md
+80
View File
@@ -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 <encode-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.
scripts/erp-tune-serve/merge_lora.py
Executable
+110
View File
@@ -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())
scripts/erp-tune-serve/quant_nvfp4a16.py
Executable
+190
View File
@@ -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())