#!/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, os, 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) # llmcompressor 0.12 introspects attention structure off the TOP-LEVEL config # (to size kv-cache/quant params). Qwen3_5 keeps num_attention_heads etc. under # text_config, and transformers 5.10 no longer delegates the top-level lookup, # so oneshot raises "Cannot determine num_attention_heads from config". Promote # them from the authoritative text_config for the duration of quant, then # restore, so the saved config keeps its canonical text_config-only shape. # (This env moved under us since the 2026-08-15 heresy quant, where the older # transformers still delegated — same "the fight is the environment" pattern.) _promote = ("num_attention_heads", "num_key_value_heads", "hidden_size", "head_dim", "num_hidden_layers") _tc = getattr(model.config, "text_config", None) _orig = {f: getattr(model.config, f, None) for f in _promote} if _tc is not None: for f in _promote: v = getattr(_tc, f, None) if v is not None: setattr(model.config, f, v) print("promoted text_config attention fields to top-level config for oneshot: " + ", ".join(f"{f}={getattr(model.config, f)}" for f in _promote), flush=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) # restore the canonical config shape (undo the promotion above) so the saved # top-level config matches the known-good heresy output; vLLM reads text_config. for f, v in _orig.items(): try: setattr(model.config, f, v) except Exception: pass print(f"saving -> {a.out}", flush=True) model.save_pretrained(a.out, save_compressed=True) # ⚠ DO NOT `tok.save_pretrained(a.out)` — that ships a CRIPPLED tokenizer. # `load_calib` calls tok(..., truncation=True, max_length=seqlen), which # MUTATES the fast tokenizer's Rust backend truncation state in place. # save_pretrained then bakes {"truncation": {"max_length": }} into # tokenizer.json, so every prompt is silently clamped at the CALIBRATION # length forever. Latent on older transformers (it does not enforce the # check) and fatal on newer ones: a vision model dies at startup because the # dummy profiling image expands to more image tokens than the cap allows # ("Mismatch in `image` token count ... Got ids=[]"). # Save a pristine tokenizer re-read from the SOURCE instead. from transformers import AutoTokenizer as _AutoTokenizer _AutoTokenizer.from_pretrained(a.model, trust_remote_code=True).save_pretrained(a.out) # Fail loudly rather than shipping the defect again. _tj = os.path.join(a.out, "tokenizer.json") if os.path.exists(_tj): with open(_tj) as _f: _trunc = json.load(_f).get("truncation") if _trunc: raise SystemExit( f"FAILED CHECK: saved tokenizer.json carries truncation={_trunc}. " "It must be null — see quant playbook §3.14." ) print("tokenizer saved pristine (truncation=null) — verified", flush=True) print("DONE (post-steps still required: graft MTP, preprocessor_config, verify ignore)", flush=True) return 0 if __name__ == "__main__": sys.exit(main())