import torchfrom datasets import load_datasetfrom transformers import AutoProcessor, Qwen3_5ForConditionalGeneration from llmcompressor import oneshotfrom llmcompressor.modifiers.gptq import GPTQModifierfrom llmcompressor.modifiers.transform.awq import AWQModifierfrom llmcompressor.utils import load_context MODEL_ID = "Qwen/Qwen3.8-27B" # Load model.with load_context(Qwen3_5ForConditionalGeneration): model = Qwen3_5ForConditionalGeneration.from_pretrained(MODEL_ID)processor = AutoProcessor.from_pretrained(MODEL_ID) recipe = [ AWQModifier(duo_scaling="both"), GPTQModifier( targets="Linear", scheme="W4A16", ignore=[ "re:visual.*", "re:model.visual.*", r"re:.*lm_head", "re:.*embed_tokens$", r"re:.*linear_attn\.in_proj_a$", r"re:.*linear_attn\.in_proj_b$", ], kv_cache_scheme={ "num_bits": 8, "type": "float", "symmetric": True, "strategy": "tensor", "dynamic": False, "observer": "static_minmax", }, ),] NUM_CALIBRATION_SAMPLES = 512MAX_SEQUENCE_LENGTH = 4096 ds = load_dataset( "mlabonne/open-perfectblend", split=f"train[:{NUM_CALIBRATION_SAMPLES}]",)ds = ds.shuffle(seed=42) ROLE_MAP = {"human": "user", "gpt": "assistant"} def preprocess_function(example): messages = [ { "role": ROLE_MAP.get(msg["from"], msg["from"]), "content": [{"type": "text", "text": msg["value"]}], } for msg in example["conversations"] ] return processor.apply_chat_template( messages, tokenize=True, return_dict=True, add_generation_prompt=False, processor_kwargs={ "return_tensors": "pt", "padding": False, "truncation": True, "max_length": MAX_SEQUENCE_LENGTH, "add_special_tokens": False, }, ) ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names) def data_collator(batch): assert len(batch) == 1 return {key: torch.tensor(value) for key, value in batch[0].items()} # Apply quantization.oneshot( model=model, recipe=recipe, dataset=ds, max_seq_length=MAX_SEQUENCE_LENGTH, num_calibration_samples=NUM_CALIBRATION_SAMPLES, moe_calibrate_all_experts=True, data_collator=data_collator,) # Save to disk in compressed-tensors format.SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-INT4"model.save_pretrained(SAVE_DIR)processor.save_pretrained(SAVE_DIR)