Views
No views yet
pip install llmcompressor1from datasets import load_dataset
2from transformers import AutoTokenizer
3from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot
4
5MODEL_ID = "mistralai/Mistral-Large-Instruct-2407"
6model = SparseAutoModelForCausalLM.from_pretrained(
7 MODEL_ID,
8 device_map="auto",
9 torch_dtype="auto",
10)
11tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
12
13# Select calibration dataset.
14DATASET_ID = "HuggingFaceH4/ultrachat_200k" # Or use your own dataset
15DATASET_SPLIT = "train_sft"
16
17# You can increase the the number of samples to increase accuracy
18NUM_CALIBRATION_SAMPLES = 512
19MAX_SEQUENCE_LENGTH = 2048
20
21ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
22ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
23
24
25def process_and_tokenize(example):
26 text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
27 return tokenizer(
28 text,
29 padding=False,
30 max_length=MAX_SEQUENCE_LENGTH,
31 truncation=True,
32 add_special_tokens=False,
33 )
34
35ds = ds.map(process_and_tokenize, remove_columns=ds.column_names)
36
37# Configure the quantization algorithm and scheme.
38# In this case, we:
39# * quantize the weights to fp8 with per-tensor scales
40# * quantize the activations to fp8 with per-tensor scales
41# * quantize the kv cache to fp8 with per-tensor scales
42recipe = """
43quant_stage:
44 quant_modifiers:
45 QuantizationModifier:
46 ignore: ["lm_head"]
47 config_groups:
48 group_0:
49 weights:
50 num_bits: 8
51 type: float
52 strategy: tensor
53 dynamic: false
54 symmetric: true
55 input_activations:
56 num_bits: 8
57 type: float
58 strategy: tensor
59 dynamic: false
60 symmetric: true
61 targets: ["Linear"]
62 kv_cache_scheme:
63 num_bits: 8
64 type: float
65 strategy: tensor
66 dynamic: false
67 symmetric: true
68"""
69
70# Apply algorithms.
71oneshot(
72 model=model,
73 dataset=ds,
74 recipe=recipe,
75 max_seq_length=MAX_SEQUENCE_LENGTH,
76 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
77)
78
79# Save to disk compressed.
80SAVE_DIR = "./Mistral-Large-Instruct-2407-FP8"
81model.save_pretrained(SAVE_DIR, save_compressed=True)
82tokenizer.save_pretrained(SAVE_DIR)