#!/usr/bin/env python3 """Qwen2.5-VL-3B: GPTQ W4/G256 text + GPTQ W8/channel vision linears. This is an experimental exception to the workspace VLM default of vision RTN. Both component recipes are calibrated using deterministic Flickr30K image-text examples so GPTQ observes the visual and language activations. """ import argparse import importlib.metadata import platform import shutil from pathlib import Path import torch from llmcompressor import oneshot from llmcompressor.args import DatasetArguments from llmcompressor.modifiers.quantization import GPTQModifier from llmcompressor.transformers.data.flickr_30k import Flickr30K from transformers import AutoModelForImageTextToText, AutoProcessor TOKENIZER_ASSET_NAMES = ( "added_tokens.json", "chat_template.jinja", "chat_template.json", "merges.txt", "sentencepiece.bpe.model", "special_tokens_map.json", "spiece.model", "tiktoken.model", "tokenizer.json", "tokenizer.model", "tokenizer_config.json", "vocab.json", "vocab.txt", ) def copy_original_tokenizer(source_dir, output_dir): """Replace generated tokenizer assets with the original files byte-for-byte.""" source_dir = Path(source_dir) output_dir = Path(output_dir) if not (source_dir / "tokenizer.json").is_file(): raise FileNotFoundError(f"Original tokenizer.json not found in {source_dir}") for name in TOKENIZER_ASSET_NAMES: source = source_dir / name destination = output_dir / name if source.is_file(): shutil.copy2(source, destination) elif destination.exists(): destination.unlink() TEXT_BITS = 4 TEXT_GROUP_SIZE = 256 VISION_BITS = 8 NUM_CALIBRATION_SAMPLES = 512 MAX_SEQ_LENGTH = 2048 BATCH_SIZE = 1 CALIBRATION_SPLIT = f"test[:{NUM_CALIBRATION_SAMPLES}]" CALIBRATION_DATASET = "lmms-lab/flickr30k" def package_version(*names: str) -> str: for name in names: try: return importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: pass return "not-installed" def create_calibration_dataset(processor): """Create deterministic image-text calibration inputs through the processor.""" args = DatasetArguments( max_seq_length=MAX_SEQ_LENGTH, num_calibration_samples=NUM_CALIBRATION_SAMPLES, batch_size=BATCH_SIZE, pad_to_max_length=False, concatenate_data=False, shuffle_calibration_samples=False, preprocessing_num_workers=1, ) # GPTQ only consumes forward activations. LLM Compressor 0.12.0 cannot # add/mask labels for samples that include pixel_values, so leave labels # out rather than taking the text-only dataset path. return Flickr30K(args, split=CALIBRATION_SPLIT, processor=processor)( add_labels=False ) def collate_qwen_vl(features): """Collate one Qwen VL sample without adding a vision batch dimension. The generic LLM Compressor collator stacks every field. Qwen passes `pixel_values` as one flattened patch sequence and `image_grid_thw` as `[images, 3]`; stacking a single example changes the latter to `[1, 1, 3]` and causes the vision position-id computation to fail. """ if len(features) != 1: raise ValueError("Qwen2.5-VL calibration requires BATCH_SIZE=1.") sample = features[0] batch = {} for name, value in sample.items(): tensor = torch.as_tensor(value) if name in {"image_grid_thw", "video_grid_thw"}: batch[name] = tensor.reshape(-1, 3) elif name in {"pixel_values", "pixel_values_videos"}: # Qwen vision forward expects flattened patches, not [batch, patches, dim]. batch[name] = tensor.squeeze(0) if tensor.ndim == 3 and tensor.shape[0] == 1 else tensor else: batch[name] = tensor.unsqueeze(0) if tensor.ndim == 1 else tensor if not getattr(collate_qwen_vl, "_logged", False): print( "Calibration batch shapes: " + ", ".join(f"{name}={tuple(value.shape)}" for name, value in batch.items()), flush=True, ) collate_qwen_vl._logged = True return batch def resolve_targets(model): """Return mutually exclusive text and Qwen `model.visual` Linear targets.""" all_linears = [ (name, module) for name, module in model.named_modules() if isinstance(module, torch.nn.Linear) ] vision_linears = [name for name, _ in all_linears if name.startswith("model.visual.")] text_linears = [name for name, _ in all_linears if not name.startswith("model.visual.")] if not vision_linears: raise ValueError("No Qwen visual Linear modules matched prefix 'model.visual.'.") if not text_linears or "lm_head" not in text_linears: raise ValueError("Expected non-visual decoder Linear modules and lm_head.") if set(vision_linears) & set(text_linears): raise ValueError("Vision and text GPTQ targets overlap.") print(f"Text GPTQ INT4/G256 Linear targets: {len(text_linears)}") print(f"Vision GPTQ INT8/channel Linear targets: {len(vision_linears)}") print(f"Vision targets (first 20): {vision_linears[:20]}") return text_linears, vision_linears def write_recipe(output_dir: Path, text_targets: list[str], vision_targets: list[str]): """Save the fully resolved effective recipe with explicit component targets.""" def target_lines(targets: list[str]) -> str: return "".join(f" - {target}\n" for target in targets) output_dir.joinpath("recipe.yaml").write_text( "# Effective Qwen2.5-VL-3B quantization recipe\n" "# Calibration: lmms-lab/flickr30k test[:512], deterministic order\n" "text_gptq_int4_g256:\n" " method: GPTQ\n" " weights:\n" " num_bits: 4\n" " type: int\n" " symmetric: true\n" " strategy: group\n" " group_size: 256\n" " targets:\n" + target_lines(text_targets) + "vision_gptq_int8_channel:\n" " method: GPTQ\n" " exception_to_default_vlm_policy: true\n" " weights:\n" " num_bits: 8\n" " type: int\n" " symmetric: true\n" " strategy: channel\n" " targets:\n" + target_lines(vision_targets) ) def validate_scales(output_dir: Path): from safetensors.torch import safe_open invalid = [] for shard in output_dir.glob("*.safetensors"): with safe_open(shard, framework="pt", device="cpu") as tensors: for name in tensors.keys(): if "scale" not in name.lower(): continue value = tensors.get_tensor(name) if value.is_floating_point() and not torch.isfinite(value.float()).all(): invalid.append(f"{shard.name}:{name}") if invalid: raise ValueError(f"Non-finite saved scale tensors: {invalid[:20]}") def save_metadata(model_path: Path, output_dir: Path): for source in model_path.iterdir(): if source.is_file() and source.suffix in {".json", ".txt"} and not source.name.endswith(".index.json"): destination = output_dir / source.name if not destination.exists(): shutil.copy2(source, destination) shutil.copy2(Path(__file__).resolve(), output_dir / "quantize.py") output_dir.joinpath("versions.txt").write_text( "\n".join( ( f"Python: {platform.python_version()}", f"torch: {torch.__version__}", f"CUDA: {torch.version.cuda or 'not-built'}", f"transformers: {package_version('transformers')}", f"llm-compressor: {package_version('llmcompressor', 'llm-compressor')}", f"auto-round: {package_version('auto-round', 'auto_round')}", f"compressed-tensors: {package_version('compressed-tensors')}", ) ) + "\n" ) def main(): parser = argparse.ArgumentParser() parser.add_argument("--model-path", required=True) parser.add_argument("--output-dir", required=True) args = parser.parse_args() model_path = Path(args.model_path) output_dir = Path(args.output_dir) if output_dir.exists() and any(output_dir.iterdir()): raise FileExistsError(f"Refusing to overwrite non-empty output directory: {output_dir}") model = AutoModelForImageTextToText.from_pretrained( model_path, device_map="auto", torch_dtype="auto", trust_remote_code=False, ) if model.config.model_type != "qwen2_5_vl": raise ValueError(f"Expected qwen2_5_vl, received {model.config.model_type!r}.") processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=False) text_targets, vision_targets = resolve_targets(model) calibration_dataset = create_calibration_dataset(processor) recipe = [ GPTQModifier( targets=text_targets, block_size=128, dampening_frac=0.01, actorder="static", offload_hessians=False, config_groups={ "text_gptq_int4_g256": { "targets": text_targets, "weights": { "num_bits": TEXT_BITS, "type": "int", "symmetric": True, "group_size": TEXT_GROUP_SIZE, "strategy": "group", "dynamic": False, }, } }, ), GPTQModifier( targets=vision_targets, block_size=128, dampening_frac=0.01, actorder="static", offload_hessians=False, config_groups={ "vision_gptq_int8_channel": { "targets": vision_targets, "weights": { "num_bits": VISION_BITS, "type": "int", "symmetric": True, "strategy": "channel", "dynamic": False, }, } }, ), ] oneshot( model=model, dataset=calibration_dataset, recipe=recipe, # Qwen2.5-VL's dynamic image-grid path is not compatible with the # LLM Compressor 0.12 sequential FX subgraph pipeline. Basic keeps # the processor-produced multimodal batch intact for each full forward. pipeline="basic", data_collator=collate_qwen_vl, max_seq_length=MAX_SEQ_LENGTH, num_calibration_samples=NUM_CALIBRATION_SAMPLES, shuffle_calibration_samples=False, ) output_dir.mkdir(parents=True, exist_ok=True) model.save_pretrained(output_dir, save_compressed=True) processor.save_pretrained(output_dir) copy_original_tokenizer(model_path, output_dir) write_recipe(output_dir, text_targets, vision_targets) save_metadata(model_path, output_dir) validate_scales(output_dir) print(f"Quantized model written to: {output_dir}") if __name__ == "__main__": main()