Views
No views yet

1import torch
2from compressed_tensors.offload import dispatch_model
3from datasets import load_dataset
4from transformers import AutoModelForImageTextToText, AutoProcessor
5
6from llmcompressor import oneshot
7from llmcompressor.modifiers.gptq import GPTQModifier
8
9MODEL_ID = "google/gemma-4-12B-it"
10model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype="auto")
11processor = AutoProcessor.from_pretrained(MODEL_ID)
12
13DATASET_ID = "neuralmagic/calibration"
14NUM_CALIBRATION_SAMPLES = 256
15MAX_SEQUENCE_LENGTH = 2048
16
17ds = load_dataset(DATASET_ID, name="LLM", split=f"train[:{NUM_CALIBRATION_SAMPLES}]")
18
19
20def preprocess_function(example):
21 messages = []
22 for message in example["messages"]:
23 messages.append(
24 {
25 "role": message["role"],
26 "content": [{"type": "text", "text": message["content"]}],
27 }
28 )
29
30 return processor.apply_chat_template(
31 messages,
32 return_tensors="pt",
33 padding=False,
34 truncation=True,
35 max_length=MAX_SEQUENCE_LENGTH,
36 tokenize=True,
37 add_special_tokens=False,
38 return_dict=True,
39 add_generation_prompt=False,
40 )
41
42
43ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
44
45
46def data_collator(batch):
47 assert len(batch) == 1
48 return {
49 key: (
50 torch.tensor(value)
51 if key != "pixel_values"
52 else torch.tensor(value, dtype=torch.bfloat16).squeeze(0)
53 )
54 for key, value in batch[0].items()
55 }
56
57
58recipe = GPTQModifier(
59 targets="Linear",
60 scheme="NVFP4",
61 ignore=[
62 "lm_head",
63 "re:.*embed_vision.*",
64 "re:.*embed_audio.*",
65 "re:.*vision_embedder.*",
66 ],
67)
68
69oneshot(
70 model=model,
71 recipe=recipe,
72 dataset=ds,
73 max_seq_length=MAX_SEQUENCE_LENGTH,
74 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
75 data_collator=data_collator,
76)
77
78print("\n\n")
79print("========== SAMPLE GENERATION ==============")
80dispatch_model(model)
81input_ids = torch.tensor(
82 [[
83 2, 105, 2364, 107, 818, 3282, 506, 7217, 563, 3730, 563,
84 1547, 106, 107, 105, 4368, 107
85 ]]
86).to(model.device)
87output = model.generate(
88 input_ids,
89 max_new_tokens=100,
90)
91print(processor.tokenizer.decode(output[0]))
92print("==========================================\n\n")
93
94SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4-GPTQ"
95model.save_pretrained(SAVE_DIR, save_compressed=True)
96processor.save_pretrained(SAVE_DIR)
97
98# Patch config: transformers renames checkpoint keys on load (vision_embedder ->
99# embed_vision), but save_pretrained reverts them. The ignore list in config.json
100# uses HF names (embed_vision) while safetensors keys use checkpoint names
101# (vision_embedder), so vllm can't match them. Add the checkpoint name explicitly.
102import json as _json
103_cfg_path = SAVE_DIR + "/config.json"
104with open(_cfg_path) as _f:
105 _cfg = _json.load(_f)
106_qcfg = _cfg.get("quantization_config")
107if _qcfg:
108 _ign = _qcfg.setdefault("ignore", [])
109 if "model.vision_embedder.patch_dense" not in _ign:
110 _ign.append("model.vision_embedder.patch_dense")
111 with open(_cfg_path, "w") as _f:
112 _json.dump(_cfg, _f, indent=2)
113 print("Patched config.json: added vision_embedder.patch_dense to ignore list")
114lm_eval --model vllm \
--model_args "pretrained=RedHatAI/gemma-4-12B-it-NVFP4,dtype=auto,max_model_len=$MAX_MODEL_LEN,add_bos_token=True,gpu_memory_utilization=0.85" \
--tasks gsm8k_platinum --num_fewshot 5 --apply_chat_template --batch_size auto
lm_eval --model vllm \
--model_args "pretrained=RedHatAI/gemma-4-12B-it-NVFP4,dtype=auto,max_model_len=$MAX_MODEL_LEN,add_bos_token=True,gpu_memory_utilization=0.85" \
--tasks wikitext --num_fewshot 0 --apply_chat_template --batch_size auto+---------------+------------------+--------------+---------------+----------+
| model_name | flexible-extract | strict-match | bits_per_byte | byte_ppl |
+---------------+------------------+--------------+---------------+----------+
| baseline-bf16 | 0.9082 | 0.8958 | 1.9125 | 3.7645 |
| NVFP4-RTN | 0.8892 | 0.8776 | 2.0945 | 4.2706 |
| NVFP4-iMatrix | 0.8974 | 0.8825 | 1.9855 | 3.9600 |
| *NVFP4-GPTQ* | 0.9016 | 0.8867 | 2.0704 | 4.2001 |
+---------------+------------------+--------------+---------------+----------+
Recovery
+---------------+------------------+--------------+---------------+----------+
| *NVFP4-GPTQ* | 99.27% | 98.98% | 92.37% | 89.6% |
+---------------+------------------+--------------+---------------+----------+