metadata
library_name: transformers
license: apache-2.0
base_model: Qwen/Qwen2.5-Coder-14B-Instruct
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
from llmcompressor.transformers.compression.helpers import calculate_offload_device_map
model_id = "Qwen/Qwen2.5-Coder-14B-Instruct"
model_out = "Qwen2.5-Coder-14B-Instruct.w8a8"
num_samples = 128
max_seq_len = 4096
tokenizer = AutoTokenizer.from_pretrained(model_id)
def preprocess_fn(example):
return {"text": tokenizer.apply_chat_template(example["messages"], add_generation_prompt=False, tokenize=False)}
ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
ds = ds.shuffle().select(range(num_samples))
ds = ds.map(preprocess_fn)
recipe = [
SmoothQuantModifier(smoothing_strength=0.8),
QuantizationModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"], dampening_frac=0.1)
]
device_map = calculate_offload_device_map(
model_id, reserve_for_hessians=True, num_gpus=1, torch_dtype="bfloat16"
)
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map=device_map,
torch_dtype="bfloat16",
)
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=max_seq_len,
num_calibration_samples=num_samples,
output_dir=model_out,
)