Views
No views yet
1import torch
2from compressed_tensors.offload import dispatch_model
3from compressed_tensors.quantization import preset_name_to_scheme
4from datasets import load_dataset
5from transformers import AutoModelForImageTextToText, AutoProcessor
6
7from llmcompressor import oneshot
8from llmcompressor.modifiers.quantization import QuantizationModifier
9from llmcompressor.modifiers.transform.imatrix import IMatrixGatherer
10
11MODEL_ID = "google/gemma-4-12B-it"
12model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype="auto")
13processor = AutoProcessor.from_pretrained(MODEL_ID)
14
15DATASET_ID = "neuralmagic/calibration"
16NUM_CALIBRATION_SAMPLES = 256
17MAX_SEQUENCE_LENGTH = 2048
18
19ds = load_dataset(DATASET_ID, name="LLM", split=f"train[:{NUM_CALIBRATION_SAMPLES}]")
20
21
22def preprocess_function(example):
23 messages = []
24 for message in example["messages"]:
25 messages.append(
26 {
27 "role": message["role"],
28 "content": [{"type": "text", "text": message["content"]}],
29 }
30 )
31
32 return processor.apply_chat_template(
33 messages,
34 return_tensors="pt",
35 padding=False,
36 truncation=True,
37 max_length=MAX_SEQUENCE_LENGTH,
38 tokenize=True,
39 add_special_tokens=False,
40 return_dict=True,
41 add_generation_prompt=False,
42 )
43
44
45ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
46
47
48def data_collator(batch):
49 assert len(batch) == 1
50 return {
51 key: (
52 torch.tensor(value)
53 if key != "pixel_values"
54 else torch.tensor(value, dtype=torch.bfloat16).squeeze(0)
55 )
56 for key, value in batch[0].items()
57 }
58
59
60scheme = preset_name_to_scheme("FP8_DYNAMIC", ["Linear"])
61scheme.weights.observer = "imatrix_mse"
62
63recipe = [
64 IMatrixGatherer(
65 ignore=[
66 "lm_head",
67 "re:.*embed_vision.*",
68 "re:.*embed_audio.*",
69 "re:.*vision_embedder.*",
70 ],
71 ),
72 QuantizationModifier(
73 config_groups={"group_0": scheme},
74 ignore=[
75 "lm_head",
76 "re:.*embed_vision.*",
77 "re:.*embed_audio.*",
78 "re:.*vision_embedder.*",
79 ],
80 ),
81]
82
83oneshot(
84 model=model,
85 recipe=recipe,
86 dataset=ds,
87 max_seq_length=MAX_SEQUENCE_LENGTH,
88 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
89 data_collator=data_collator,
90)
91
92print("\n\n")
93print("========== SAMPLE GENERATION ==============")
94dispatch_model(model)
95input_ids = torch.tensor(
96 [[
97 2, 105, 2364, 107, 818, 3282, 506, 7217, 563, 3730, 563,
98 1547, 106, 107, 105, 4368, 107
99 ]]
100).to(model.device)
101output = model.generate(
102 input_ids,
103 max_new_tokens=100,
104)
105print(processor.tokenizer.decode(output[0]))
106print("==========================================\n\n")
107
108SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8_Dynamic-iMatrix"
109model.save_pretrained(SAVE_DIR, save_compressed=True)
110processor.save_pretrained(SAVE_DIR)
111
112# Patch config: transformers renames checkpoint keys on load (vision_embedder ->
113# embed_vision), but save_pretrained reverts them. The ignore list in config.json
114# uses HF names (embed_vision) while safetensors keys use checkpoint names
115# (vision_embedder), so vllm can't match them. Add the checkpoint name explicitly.
116import json as _json
117_cfg_path = SAVE_DIR + "/config.json"
118with open(_cfg_path) as _f:
119 _cfg = _json.load(_f)
120_qcfg = _cfg.get("quantization_config")
121if _qcfg:
122 _ign = _qcfg.setdefault("ignore", [])
123 if "model.vision_embedder.patch_dense" not in _ign:
124 _ign.append("model.vision_embedder.patch_dense")
125 with open(_cfg_path, "w") as _f:
126 _json.dump(_cfg, _f, indent=2)
127 print("Patched config.json: added vision_embedder.patch_dense to ignore list")
128
129lm_eval --model vllm \
--model_args "pretrained=RedHatAI/gemma-4-12B-it-FP8_Dynamic,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-FP8_Dynamic,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 |
| FP8-RTN | 0.9115 | 0.8999 | 1.9368 | 3.8285 |
| *FP8-iMatrix* | 0.9198 | 0.9032 | 1.9056 | 3.7465 |
| FP8-GPTQ | 0.9098 | 0.8950 | 1.9357 | 3.8257 |
+---------------+------------------+--------------+----------------+----------+
Recovery
+---------------+------------------+--------------+---------------+----------+
|*NVFP4-iMatrix*| 100.17% | 100.83% | 100.36% | 100.48% |
+---------------+------------------+--------------+---------------+----------+