| import os
|
| import shutil
|
| import torch
|
| from datasets import load_dataset
|
| from transformers import AutoTokenizer
|
|
|
| os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
| os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
| if hasattr(torch, "accelerator") and not hasattr(torch.accelerator, "get_memory_info"):
|
| torch.accelerator.get_memory_info = lambda idx=0: torch.cuda.mem_get_info(idx)
|
|
|
| try:
|
| from transformers import AutoModelForImageTextToText
|
| AutoModelClass = AutoModelForImageTextToText
|
| except ImportError:
|
| from transformers import AutoModelForCausalLM
|
| AutoModelClass = AutoModelForCausalLM
|
|
|
| from llmcompressor import oneshot
|
| from llmcompressor.modifiers.awq import AWQModifier
|
|
|
|
|
| MODEL_ID = "/root/autodl-tmp/Huihui-Qwen3.5-9B-abliterated"
|
| SAVE_DIR = "/root/autodl-tmp/Huihui-Qwen3.5-9B-abliterated-AWQ-W4A16-vLLM"
|
|
|
| DATASET_DIR = "/root/autodl-tmp/pile-val-backup"
|
|
|
| MAX_SEQ_LENGTH = 2048
|
| NUM_CALIBRATION_SAMPLES = 128
|
|
|
|
|
| recipe = [
|
| AWQModifier(
|
| config_groups={
|
| "group_1": {
|
| "targets": ["Linear"],
|
| "weights": {
|
| "num_bits": 4,
|
| "type": "int",
|
| "symmetric": True,
|
| "group_size": 32,
|
| "strategy": "group",
|
| "observer": "mse",
|
| "dynamic": False,
|
| },
|
| "input_activations": None,
|
| "output_activations": None,
|
| }
|
| },
|
| targets=["Linear"],
|
| ignore=[
|
| "re:.*embed_tokens",
|
| "re:.*linear_attn[.]in_proj_a",
|
| "re:.*linear_attn[.]in_proj_b",
|
| "re:model[.]visual.*",
|
| "re:mtp.*",
|
| "lm_head",
|
| ],
|
| mappings=[
|
| {
|
| "smooth_layer": r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]input_layernorm",
|
| "balance_layers": [
|
| r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]q_proj",
|
| r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]k_proj",
|
| r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]v_proj",
|
| ],
|
| },
|
| {
|
| "smooth_layer": r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]v_proj",
|
| "balance_layers": [
|
| r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]o_proj",
|
| ],
|
| },
|
| {
|
| "smooth_layer": r"re:model.*post_attention_layernorm",
|
| "balance_layers": [
|
| r"re:model.*mlp[.]gate_proj",
|
| r"re:model.*mlp[.]up_proj",
|
| ],
|
| },
|
| ],
|
| duo_scaling=True,
|
| n_grid=20,
|
| )
|
| ]
|
|
|
|
|
| print("Loading model...")
|
| model = AutoModelClass.from_pretrained(
|
| MODEL_ID,
|
| dtype="auto",
|
| device_map="auto",
|
| trust_remote_code=True,
|
| )
|
|
|
| print("Loading tokenizer...")
|
| tokenizer = AutoTokenizer.from_pretrained(
|
| MODEL_ID,
|
| trust_remote_code=True,
|
| )
|
|
|
| print("Loading calibration dataset...")
|
| dataset = load_dataset(
|
| DATASET_DIR,
|
| split=f"validation[:{NUM_CALIBRATION_SAMPLES}]",
|
| )
|
|
|
| print("Starting AWQ quantization...")
|
| oneshot(
|
| model=model,
|
| tokenizer=tokenizer,
|
| dataset=dataset,
|
| recipe=recipe,
|
| output_dir=SAVE_DIR,
|
| max_seq_length=MAX_SEQ_LENGTH,
|
| num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
| )
|
|
|
|
|
| extra_files = [
|
| "preprocessor_config.json",
|
| "processor_config.json",
|
| "video_preprocessor_config.json",
|
| "chat_template.jinja",
|
| "generation_config.json",
|
| "tokenizer.json",
|
| "tokenizer_config.json",
|
| "vocab.json",
|
| "merges.txt",
|
| ]
|
|
|
| for name in extra_files:
|
| src = os.path.join(MODEL_ID, name)
|
| dst = os.path.join(SAVE_DIR, name)
|
| if os.path.exists(src) and not os.path.exists(dst):
|
| shutil.copy2(src, dst)
|
|
|
| print(f"Saved vLLM-compatible AWQ W4A16 model to: {SAVE_DIR}") |
