feat(gen-seat): mixed NVFP4+FP8 requant — +18% decode at equal MTP acceptance
Re-quantizes the fleet `gen` seat from weight-only NVFP4A16 to a mixed-precision build: NVFP4 W4A4 for layers 0-55 MLPs, FP8 W8A8 for the attention projections / linear_attn / lm_head / layers 56-63 MLPs, FP8 KV cache. Replicates the scheme of unsloth/Qwen3.8-27B-NVFP4 on the abliterated weights. The queued task named this "W4A8" (NVFP4 weights + FP8 activations). That checkpoint cannot be served: vLLM 0.24's compressed-tensors dispatcher (compressed_tensors.py:704-713) accepts NVFP4 weights with either no input quantization (W4A16, which forces the Marlin kernel) or NVFP4 input quantization (W4A4) -- anything else, FP8 included, raises ValueError at load. CompressedTensorsW4A8Fp8 is INT4 weights gated on an exact-sm90 check, so it is closed on Blackwell twice over. The ~20% intuition was correct; the scheme name was not. Getting FP8 into the mix has to be done per-layer-group. Established the gain before spending GPU time: unsloth's build was already on-box, so serving it as a probe measured +19.1% over our seat at identical MTP acceptance -- a kernel-level result, no requant needed to learn it. Measured, cache-busted, bs=1: decode 80.12 -> 94.53 tok/s (+18.0%) MTP acceptance 47.8% -> 47.7% (unchanged) perplexity (n=6) 6.941 -> 7.059 (+1.7%) abliteration 4/4 -> 4/4 (preserved) weights on disk 27.7 -> 22.5 GB (-19%) Surface test green on the live seat: plain chat, vision, tool calling, thinking split, 36K-token needle retrieval, streaming. All 7 LiteLLM aliases verified routing. GEN_GPU_MEM_UTIL 0.45 -> 0.43: the new weights are 5.2 GB smaller, and at 0.45 the seat absorbed that slack as KV, leaving meromero-charrp 0.18 GiB short of its budget on the shared GPU0 -- it crash-looped. Handing the space back leaves gen 422K tokens of KV (1.6x its 262K context) and both seats co-resident at 89.8/97.9 GB. Also records two measured negatives so they are not re-chased: GEN_SPEC_TOKENS is already optimal at 3 (swept 2/3/4/5 -> 77.1/80.1/78.7/ 75.9 tok/s), and vLLM's prompt_logprobs are ~uniform while speculative decoding is on, so perplexity must be measured with spec off. Pipeline, acceptance harness and raw measurements land in services/gen-seat-mixed-quant/. Rollback is one .env line; the previous build is untouched at /tank/aimodels/qwen38-27b-uncensored-nvfp4.
This commit is contained in:
@@ -0,0 +1,145 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Mixed-precision NVFP4(W4A4 MLP) + FP8(W8A8 attn) quant of Qwen3.8-27B-Uncensored.
|
||||
|
||||
Replicates the scheme of `unsloth/Qwen3.8-27B-NVFP4` (verified on-box to run
|
||||
+19.1% faster than our weight-only NVFP4A16/Marlin build at identical MTP
|
||||
acceptance), applied to the abliterated uncensored weights.
|
||||
|
||||
WHY THIS SHAPE, not the "NVFP4 weights + FP8 activations" W4A8 idea:
|
||||
vLLM 0.24's compressed-tensors dispatcher (compressed_tensors.py:704-713) allows
|
||||
NVFP4 weights with EXACTLY two activation options -- None (=W4A16, Marlin) or
|
||||
NVFP4 (=W4A4). Anything else, FP8 included, raises
|
||||
ValueError: For NVFP4 weights, input quantization must also be NVFP4 format
|
||||
So a literal W4A8-on-NVFP4 checkpoint cannot load. The servable way to get FP8
|
||||
into the mix is per-layer-group: NVFP4 W4A4 for the bulk MLPs, FP8 W8A8 for the
|
||||
attention projections / linear_attn / lm_head / last-8-layer MLPs.
|
||||
|
||||
Groups (byte-for-byte the unsloth recipe):
|
||||
group_0 FP8 W8A8 channel weights (static) + per-token dynamic activations
|
||||
-> self_attn q/k/v/o, linear_attn in_proj_qkv/in_proj_z/out_proj,
|
||||
lm_head, and layers 56-63 MLPs (late layers are accuracy-sensitive)
|
||||
group_1 NVFP4 W4A4 tensor_group gsize16, fp8 scales, imatrix_mse weights
|
||||
-> layers 0-55 MLP gate/up/down
|
||||
kv_cache FP8 static tensor
|
||||
|
||||
Targets are made EXPLICITLY non-overlapping (group_1 enumerates layers 0-55)
|
||||
rather than relying on group-precedence to resolve the 56-63 collision.
|
||||
|
||||
Kept out entirely: vision tower, linear_attn norm/in_proj_a/in_proj_b, and MTP.
|
||||
MTP MUST stay in `ignore` -- otherwise vLLM loads the grafted bf16 MTP head as
|
||||
quantized and it comes up uninitialised (0% acceptance). That bug cost two prior
|
||||
rounds; do not remove `re:^mtp.*`.
|
||||
"""
|
||||
import argparse, json, sys
|
||||
|
||||
# --- group_0: FP8 W8A8 -------------------------------------------------------
|
||||
G0_TARGETS = [
|
||||
r"re:.*self_attn\.(q|k|v|o)_proj$",
|
||||
r"re:.*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$",
|
||||
r"re:.*lm_head",
|
||||
r"re:.*layers\.(56|57|58|59|60|61|62|63)\.mlp\.(gate|up|down)_proj$",
|
||||
]
|
||||
# --- group_1: NVFP4 W4A4, layers 0-55 only (0-9 | 10-49 | 50-55) -------------
|
||||
G1_TARGETS = [
|
||||
r"re:.*layers\.([0-9]|[1-4][0-9]|5[0-5])\.mlp\.(gate|up|down)_proj$",
|
||||
]
|
||||
IGNORE = [
|
||||
r"re:.*visual.*",
|
||||
r"re:.*linear_attn\.(norm|in_proj_a|in_proj_b)$",
|
||||
r"re:^mtp.*",
|
||||
]
|
||||
|
||||
|
||||
def build_recipe():
|
||||
from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme
|
||||
from llmcompressor.modifiers.quantization import QuantizationModifier
|
||||
|
||||
g0 = QuantizationScheme(
|
||||
targets=G0_TARGETS,
|
||||
weights=QuantizationArgs(num_bits=8, type="float", strategy="channel",
|
||||
symmetric=True, dynamic=False,
|
||||
observer="memoryless_minmax"),
|
||||
input_activations=QuantizationArgs(num_bits=8, type="float", strategy="token",
|
||||
symmetric=True, dynamic=True),
|
||||
)
|
||||
g1 = QuantizationScheme(
|
||||
targets=G1_TARGETS,
|
||||
weights=QuantizationArgs(num_bits=4, type="float", strategy="tensor_group",
|
||||
group_size=16, symmetric=True, dynamic=False,
|
||||
observer="imatrix_mse", actorder="static",
|
||||
scale_dtype="torch.float8_e4m3fn"),
|
||||
input_activations=QuantizationArgs(num_bits=4, type="float", strategy="tensor_group",
|
||||
group_size=16, symmetric=True, dynamic="local",
|
||||
observer="static_minmax",
|
||||
scale_dtype="torch.float8_e4m3fn"),
|
||||
)
|
||||
kv = QuantizationArgs(num_bits=8, type="float", strategy="tensor",
|
||||
symmetric=True, dynamic=False, observer="static_minmax")
|
||||
return QuantizationModifier(
|
||||
config_groups={"group_0": g0, "group_1": g1},
|
||||
ignore=IGNORE,
|
||||
kv_cache_scheme=kv,
|
||||
)
|
||||
|
||||
|
||||
def load_calib(path, tok, n, seqlen):
|
||||
from datasets import Dataset
|
||||
rows = []
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
msgs = json.loads(line).get("messages")
|
||||
if not msgs:
|
||||
continue
|
||||
try:
|
||||
txt = tok.apply_chat_template(msgs, tokenize=False)
|
||||
except Exception:
|
||||
continue
|
||||
rows.append({"text": txt})
|
||||
if len(rows) >= n:
|
||||
break
|
||||
print(f"calibration rows: {len(rows)}", flush=True)
|
||||
ds = Dataset.from_list(rows)
|
||||
|
||||
def tokenize(b):
|
||||
return tok(b["text"], truncation=True, max_length=seqlen, add_special_tokens=False)
|
||||
|
||||
return ds.map(tokenize, remove_columns=["text"])
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--calib", required=True)
|
||||
ap.add_argument("--out", required=True)
|
||||
ap.add_argument("--num-samples", type=int, default=256)
|
||||
ap.add_argument("--seqlen", type=int, default=2048)
|
||||
a = ap.parse_args()
|
||||
|
||||
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration
|
||||
from llmcompressor import oneshot
|
||||
|
||||
print(f"loading (wrapper class) {a.model}", flush=True)
|
||||
tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
|
||||
model = Qwen3_5ForConditionalGeneration.from_pretrained(
|
||||
a.model, torch_dtype="auto", device_map=None, trust_remote_code=True)
|
||||
|
||||
ds = load_calib(a.calib, tok, a.num_samples, a.seqlen)
|
||||
recipe = build_recipe()
|
||||
print("oneshot: NVFP4 W4A4 (L0-55 MLP) + FP8 W8A8 (attn/linear_attn/lm_head/L56-63 MLP) "
|
||||
"+ FP8 KV; vision/linear_attn-norms/MTP ignored", flush=True)
|
||||
oneshot(model=model, dataset=ds, recipe=recipe,
|
||||
num_calibration_samples=len(ds), max_seq_length=a.seqlen)
|
||||
|
||||
print(f"saving -> {a.out}", flush=True)
|
||||
model.save_pretrained(a.out, save_compressed=True)
|
||||
tok.save_pretrained(a.out)
|
||||
print("DONE (post-steps still required: graft MTP, preprocessor_config, verify ignore)",
|
||||
flush=True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user