Views
No views yet
VLLM_MARLIN_INPUT_DTYPE=int8.| Tasks |Version| Filter |n-shot| Metric | |Value| |Stderr|
|---------------------|-------|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k | 3|flexible-extract| 5|exact_match|↑ |0.100|± |0.0190|
| | |strict-match | 5|exact_match|↑ |0.072|± |0.0164|
|kormedmcqa | 2|none | |exact_match|↑ |0.755|± |0.0134|
| - kormedmcqa_dentist| 2|none | 5|exact_match|↑ |0.704|± |0.0289|
| - kormedmcqa_doctor | 2|none | 5|exact_match|↑ |0.644|± |0.0303|
| - kormedmcqa_nurse | 2|none | 5|exact_match|↑ |0.844|± |0.0230|
| - kormedmcqa_pharm | 2|none | 5|exact_match|↑ |0.828|± |0.0239|
|medmcqa |Yaml |none | 5|acc |↑ |0.676|± |0.0297|
| | |none | 5|acc_norm |↑ |0.676|± |0.0297|
| Groups |Version|Filter|n-shot| Metric | |Value| |Stderr|
|----------|------:|------|------|-----------|---|----:|---|-----:|
|kormedmcqa| 2|none | |exact_match|↑ |0.755|± |0.0134|| Tasks |Version| Filter |n-shot| Metric | |Value| |Stderr|
|---------------------|-------|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k | 3|flexible-extract| 5|exact_match|↑ |0.088|± |0.0180|
| | |strict-match | 5|exact_match|↑ |0.056|± |0.0146|
|kormedmcqa | 2|none | |exact_match|↑ |0.733|± |0.0137|
| - kormedmcqa_dentist| 2|none | 5|exact_match|↑ |0.672|± |0.0298|
| - kormedmcqa_doctor | 2|none | 5|exact_match|↑ |0.620|± |0.0308|
| - kormedmcqa_nurse | 2|none | 5|exact_match|↑ |0.828|± |0.0239|
| - kormedmcqa_pharm | 2|none | 5|exact_match|↑ |0.812|± |0.0248|
|medmcqa |Yaml |none | 5|acc |↑ |0.672|± |0.0298|
| | |none | 5|acc_norm |↑ |0.672|± |0.0298|
| Groups |Version|Filter|n-shot| Metric | |Value| |Stderr|
|----------|------:|------|------|-----------|---|----:|---|-----:|
|kormedmcqa| 2|none | |exact_match|↑ |0.733|± |0.0137|1import torch
2import random
3import base64
4from io import BytesIO
5from datasets import load_dataset, Dataset, VerificationMode
6from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
7from qwen_vl_utils import process_vision_info
8
9from llmcompressor import oneshot
10from llmcompressor.modifiers.quantization import QuantizationModifier
11from llmcompressor.utils import dispatch_for_generation
12
13# ==========================================
14# 1. Configuration
15# ==========================================
16MODEL_ID = "Qwen/Qwen3-VL-30B-A3B-Thinking"
17DATASET_TEXT_ID = "neuralmagic/calibration"
18DATASET_IMG_ID = "lmms-lab/flickr30k"
19
20# Total samples for calibration
21TOTAL_SAMPLES = 512
22# 80% Text, 20% Image
23NUM_TEXT_SAMPLES = int(TOTAL_SAMPLES * 0.8)
24NUM_IMG_SAMPLES = TOTAL_SAMPLES - NUM_TEXT_SAMPLES
25MAX_SEQUENCE_LENGTH = 1024
26
27# ==========================================
28# 2. Model Loading
29# ==========================================
30print(f"Loading {MODEL_ID}...")
31model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
32 MODEL_ID,
33 dtype="auto",
34 device_map="auto",
35 trust_remote_code=True
36)
37processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
38
39# ==========================================
40# 3. Dataset Processing (Hybrid 80/20)
41# ==========================================
42
43def process_text_sample(example):
44 """Preprocesses a text-only sample from neuralmagic/calibration."""
45 messages = []
46 for message in example["messages"]:
47 messages.append(
48 {
49 "role": message["role"],
50 "content": [{"type": "text", "text": message["content"]}],
51 }
52 )
53
54 text_inputs = processor.apply_chat_template(
55 messages,
56 return_tensors="pt",
57 padding=False,
58 truncation=True,
59 max_length=MAX_SEQUENCE_LENGTH,
60 tokenize=True,
61 add_special_tokens=False,
62 return_dict=True,
63 add_generation_prompt=False,
64 )
65
66 return {
67 "input_ids": text_inputs.input_ids[0],
68 "attention_mask": text_inputs.attention_mask[0],
69 # Explicitly set visual keys to None for text samples
70 "pixel_values": None,
71 "image_grid_thw": None,
72 "video_grid_thw": None
73 }
74
75def process_image_sample(example):
76 """Preprocesses an image sample from flickr30k."""
77 # Convert PIL image to base64 for the chat template
78 buffered = BytesIO()
79 example["image"].save(buffered, format="PNG")
80 encoded_image = base64.b64encode(buffered.getvalue())
81 encoded_image_text = encoded_image.decode("utf-8")
82 base64_qwen = f"data:image;base64,{encoded_image_text}"
83
84 # Create a generic prompt for the image
85 messages = [
86 {
87 "role": "user",
88 "content": [
89 {"type": "image", "image": base64_qwen},
90 {"type": "text", "text": "Describe this image in detail."},
91 ],
92 }
93 ]
94
95 # Process text formatting
96 text = processor.apply_chat_template(
97 messages, tokenize=False, add_generation_prompt=True
98 )
99
100 # Process visual info (extracts pixel values and grid info)
101 image_inputs, video_inputs = process_vision_info(messages)
102
103 # Tokenize and create tensors
104 inputs = processor(
105 text=[text],
106 images=image_inputs,
107 videos=video_inputs,
108 padding=False,
109 max_length=MAX_SEQUENCE_LENGTH,
110 truncation=True,
111 )
112
113 return {
114 "input_ids": torch.tensor(inputs["input_ids"][0]),
115 "attention_mask": torch.tensor(inputs["attention_mask"][0]),
116 # Qwen3 VL visual outputs
117 "pixel_values": torch.tensor(inputs["pixel_values"]),
118 "image_grid_thw": torch.tensor(inputs["image_grid_thw"]),
119 "video_grid_thw": None # Assuming no video in flickr30k
120 }
121
122print(f"Preparing datasets: {NUM_TEXT_SAMPLES} Text / {NUM_IMG_SAMPLES} Images...")
123
124# Load and process Text Dataset
125ds_text = load_dataset(DATASET_TEXT_ID, name="LLM", split=f"train[:{NUM_TEXT_SAMPLES}]")
126processed_data = []
127
128for sample in ds_text:
129 processed_data.append(process_text_sample(sample))
130
131# Load and process Image Dataset
132ds_img = load_dataset(DATASET_IMG_ID, data_files="data/test-00001-of-00009.parquet", split=f"train[:{NUM_IMG_SAMPLES}]", verification_mode=VerificationMode.NO_CHECKS)
133for sample in ds_img:
134 processed_data.append(process_image_sample(sample))
135
136random.shuffle(processed_data)
137combined_ds = Dataset.from_list(processed_data)
138
139# ==========================================
140# 4. Data Collator
141# ==========================================
142def hybrid_data_collator(batch):
143 """
144 Handles batches that might be text-only OR multimodal.
145 Removes None values (which represent missing visual data in text samples).
146 """
147 assert len(batch) == 1, "Batch size must be 1 for oneshot calibration"
148 sample = batch[0]
149
150 batch_out = {}
151
152 # Handle Input IDs and Mask (Always present)
153 batch_out["input_ids"] = torch.tensor(sample["input_ids"]).unsqueeze(0)
154 batch_out["attention_mask"] = torch.tensor(sample["attention_mask"]).unsqueeze(0)
155
156 # Handle Visuals (Only present if not None)
157 if sample.get("pixel_values") is not None:
158 # Convert list back to tensor if dataset conversion made them lists
159 batch_out["pixel_values"] = torch.tensor(sample["pixel_values"])
160
161 # Qwen3 usually requires bfloat16 for pixel values
162 batch_out["pixel_values"] = batch_out["pixel_values"].to(dtype=torch.bfloat16)
163
164 if sample.get("image_grid_thw") is not None:
165 batch_out["image_grid_thw"] = torch.tensor(sample["image_grid_thw"])
166
167 if sample.get("video_grid_thw") is not None:
168 batch_out["video_grid_thw"] = torch.tensor(sample["video_grid_thw"])
169
170 return batch_out
171
172# ==========================================
173# 5. Quantization
174# ==========================================
175recipe = QuantizationModifier(
176 targets="Linear",
177 scheme="W4A16",
178 ignore=[
179 "re:.*embed_tokens", # Don't mess with token embedding space
180# "re:.*self_attn.*",
181 "re:.*lm_head",
182 "re:visual.*",
183 "re:model.visual.*",
184 "re:.*mlp.gate$", # Gate is crucial for MoE
185 ],
186)
187
188print("Starting One-Shot Calibration...")
189SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-W4A16"
190oneshot(
191 model=model,
192 recipe=recipe,
193 max_seq_length=MAX_SEQUENCE_LENGTH,
194 num_calibration_samples=TOTAL_SAMPLES,
195 dataset=combined_ds,
196 data_collator=hybrid_data_collator,
197 moe_calibrate_all_experts=True,
198 output_dir=SAVE_DIR
199)