import jsonimport os import torchfrom datasets import load_datasetfrom transformers import AutoModelForImageTextToText, AutoTokenizer from llmcompressor import oneshotfrom llmcompressor.modifiers.quantization.gptq import GPTQModifier MODEL_ID = "Qwen/Qwen3.6-35B-A3B"OUTPUT_DIR = "./Qwen3.6-35B-A3B-w4a16-llmcompressor-v0.12.0"CALIB_SIZE = 128MAX_SEQ_LENGTH = 2048 IGNORE = [ "lm_head", r"re:.*\.router$", r"re:.*\.router\..*", r"re:.*\.gate$", r"re:.*\.mlp\.gate$", "shared_expert_gate", r"re:.*\.shared_expert_gate(\.|$)", r"re:.*\.visual(\.|$)", r"re:.*\.visual\..*",] # Step 1: Load the BF16 model with the multimodal wrapper so vision_config and# the ForConditionalGeneration architecture are preserved on save.model = AutoModelForImageTextToText.from_pretrained( MODEL_ID, dtype=torch.bfloat16, device_map="cpu", trust_remote_code=True,)reference_config = model.configtokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) # Step 2: Load calibration data. GPTQ is data-driven: it needs real activations# to build the per-layer Hessians used to compensate the rounding error.ds = load_dataset("HuggingFaceH4/ultrachat_200k", split=f"train_sft[:{CALIB_SIZE}]")ds = ds.map( lambda ex: {"text": "\n".join(m["content"] for m in ex["messages"] if m.get("content"))}, remove_columns=ds.column_names,)if not getattr(tokenizer, "pad_token", None): tokenizer.pad_token = tokenizer.eos_tokencalib_ds = ds.map( lambda ex: tokenizer( ex["text"], truncation=True, max_length=MAX_SEQ_LENGTH, add_special_tokens=False ), remove_columns=["text"],) # Step 3: Define the W4A16 GPTQ recipe and one-shot quantize. Passing# processor=MODEL_ID saves the processor configs alongside the model.recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=IGNORE)oneshot( model=model, dataset=calib_ds, recipe=recipe, max_seq_length=MAX_SEQ_LENGTH, processor=MODEL_ID, output_dir=OUTPUT_DIR, save_compressed=True, trust_remote_code_model=True,) # Step 4: Patch config.json so architectures/vision_config stay intact for vLLM.config_path = os.path.join(OUTPUT_DIR, "config.json")with open(config_path, encoding="utf-8") as f: saved = json.load(f)full = reference_config.to_dict()for key in ( "architectures", "model_type", "text_config", "vision_config", "vision_start_token_id", "vision_end_token_id", "image_token_id", "video_token_id",): if full.get(key) is not None: saved[key] = full[key]with open(config_path, "w", encoding="utf-8") as f: json.dump(saved, f, indent=2) f.write("\n") # Smoke testinputs = tokenizer("What are we having for dinner?", return_tensors="pt")output = model.generate(**inputs, max_new_tokens=30)print(tokenizer.decode(output[0], skip_special_tokens=True))