esh-pfi-infrastructure/services/gen-seat-mixed-quant/quant_mixed_nvfp4.py

#!/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": <seqlen>}} 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=[<cap-1>]").
    # 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())