Views
No views yet
1from transformers import AutoTokenizer
2from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot
3from llmcompressor.modifiers.quantization import GPTQModifier
4from compressed_tensors.quantization import QuantizationArgs, QuantizationType, QuantizationStrategy
5from datasets import load_dataset
6import random
7
8model_id = "HuggingFaceTB/SmolLM-135M-Instruct"
9
10
11num_samples = 512
12max_seq_len = 4096
13
14tokenizer = AutoTokenizer.from_pretrained(model_id)
15
16preprocess_fn = lambda example: {"text": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n{text}".format_map(example)}
17
18dataset_name = "neuralmagic/LLM_compression_calibration"
19dataset = load_dataset(dataset_name, split="train")
20ds = dataset.shuffle().select(range(num_samples))
21ds = ds.map(preprocess_fn)
22
23examples = [
24 tokenizer(
25 example["text"], padding=False, max_length=max_seq_len, truncation=True,
26 ) for example in ds
27]
28
29# recipe = "w4a16_nohead_recipe.yaml"
30recipe = GPTQModifier(
31 targets="Linear",
32 scheme="W4A16",
33 ignore=["lm_head"],
34 dampening_frac=0.1,
35)
36
37
38model = SparseAutoModelForCausalLM.from_pretrained(
39 model_id,
40 device_map="auto",
41 trust_remote_code=True
42)
43
44print(model)
45
46oneshot(
47 model=model,
48 dataset=ds,
49 recipe=recipe,
50 max_seq_length=max_seq_len,
51 num_calibration_samples=num_samples,
52 oneshot_device="cuda:1,2,3",
53)
54
55model_name = model_id.split("/")[-1]
56
57model.save_pretrained(f"{model_name}-quantized.w4a16")
58tokenizer.save_pretrained(f"{model_name}-quantized.w4a16")
59lm_eval \
--model sparseml \
--model_args pretrained=nm-testing/SmolLM-1.7B-Instruct-quantized.w4a16,dtype=bfloat16,max_legth=2048,add_bos_token=True,parallelize=True \
--tasks openllm \
--batch_size auto| Benchmark | SmolLM-135M-Instruct | SmolLM-135M-Instruct-quantized.w4a16(this model) | Recovery |
| MMLU (5-shot) | 26.220 | 25.202 | 96.12% |
| ARC Challenge (25-shot) | 29.948 | 30.034 | 100.29% |
| GSM-8K (5-shot, strict-match) | 1.289 | 1.971 | 152.91% |
| Hellaswag (10-shot) | 41.41 | 40.81 | 98.55% |
| Winogrande (5-shot) | 50.039 | 53.591 | 107.10% |
| TruthfulQA (0-shot) | 40.38 | 39.87 | 98.74% |
| Average | 31.55 | 31.91 | 101.16% |