File size: 3,130 Bytes
a58490a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """Quantize the Clef-Flash backbone with llm-compressor (0.14, transformers 5.17).
python build_data.py # writes calib.jsonl / eval.jsonl
python quantize.py --scheme FP8_DYNAMIC # no calibration needed
python quantize.py --scheme NVFP4 # calibrates on calib.jsonl
The joint head is not quantized; copy clef_head/ from the release next to the output.
"""
import argparse
import json
import shutil
import sys
from pathlib import Path
import torch
from datasets import Dataset
from huggingface_hub import snapshot_download
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.observers.helpers import FUSED_LAYER_NAMES
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
ap = argparse.ArgumentParser()
ap.add_argument("--release", default="Cloudflare/clef-flash", help="hub id or local path of the BF16 release")
ap.add_argument("--scheme", required=True, choices=["FP8_DYNAMIC", "FP8", "NVFP4"])
ap.add_argument("--calib", default="calib.jsonl")
ap.add_argument("--num-samples", type=int, default=672)
ap.add_argument("--pipeline", default="sequential")
args = ap.parse_args()
RELEASE = Path(args.release) if Path(args.release).is_dir() else Path(snapshot_download(args.release))
sys.path.insert(0, str(RELEASE))
from joint_schema_model import encode_record # noqa: E402
IGNORE = [
"lm_head", # head reads it as lexical option embeddings
"re:.*visual.*", # vision tower
"re:.*linear_attn.in_proj_a$", # GDN gate projections, 32 outputs each
"re:.*linear_attn.in_proj_b$",
]
# vLLM fuses Qwen3.5 GDN in_proj_qkv + in_proj_z into one in_proj_qkvz GEMM, so NVFP4 needs a
# shared weight global scale for them (llm-compressor only knows q/k/v and gate/up). Without
# this vLLM takes the max scale and the qkv weights come out scaled by g_qkv / g_z.
FUSED_LAYER_NAMES.append(("in_proj_qkv", "in_proj_z"))
processor = AutoProcessor.from_pretrained(RELEASE)
model = Qwen3_5ForConditionalGeneration.from_pretrained(RELEASE, dtype=torch.bfloat16)
dataset = None
if args.scheme != "FP8_DYNAMIC":
rows = [json.loads(l) for l in open(args.calib)][: args.num_samples]
ids = [list(encode_record(processor.tokenizer, r["record"], processor=processor).input_ids) for r in rows]
dataset = Dataset.from_dict({"input_ids": ids, "attention_mask": [[1] * len(x) for x in ids]})
oneshot(
model=model,
dataset=dataset,
recipe=QuantizationModifier(targets="Linear", scheme=args.scheme, ignore=IGNORE),
max_seq_length=16384,
num_calibration_samples=len(dataset) if dataset is not None else 0,
pipeline=args.pipeline if dataset is not None else "datafree",
sequential_targets=["Qwen3_5DecoderLayer"],
)
out = Path("quant") / f"clef-flash-{args.scheme.lower()}"
model.save_pretrained(out, save_compressed=True)
for f in ["tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "processor_config.json",
"generation_config.json"]:
shutil.copy(RELEASE / f, out / f)
print("saved", out)
|