Views
No views yet
1from datasets import load_dataset
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4from llmcompressor import oneshot
5from llmcompressor.modifiers.quantization import QuantizationModifier
6from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
7from llmcompressor.utils import dispatch_for_generation
8
9MODEL_ID = "meta-llama/Llama-2-7b-hf"
10
11# Load model.
12model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
13tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
14
15
16DATASET_ID = "HuggingFaceH4/ultrachat_200k"
17DATASET_SPLIT = "train_sft"
18
19NUM_CALIBRATION_SAMPLES = 512
20MAX_SEQUENCE_LENGTH = 2048
21
22# Load dataset and preprocess.
23ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
24ds = ds.shuffle(seed=42)
25
26
27def preprocess(example):
28 return {
29 "text": tokenizer.apply_chat_template(
30 example["messages"],
31 tokenize=False,
32 )
33 }
34
35
36ds = ds.map(preprocess)
37
38
39# Tokenize inputs.
40def tokenize(sample):
41 return tokenizer(
42 sample["text"],
43 padding=False,
44 max_length=MAX_SEQUENCE_LENGTH,
45 truncation=True,
46 add_special_tokens=False,
47 )
48
49
50ds = ds.map(tokenize, remove_columns=ds.column_names)
51
52smoothing_strength = 0.5
53recipe = [
54 SmoothQuantModifier(smoothing_strength=smoothing_strength),
55 QuantizationModifier(
56 ignore=["re:.*lm_head.*"],
57 config_groups={
58 "group_0": {
59 "targets": ["Linear"],
60 "weights": {
61 "num_bits": 4,
62 "type": "float",
63 "strategy": "tensor_group",
64 "group_size": 16,
65 "symmetric": True,
66 "observer": "minmax",
67 },
68 "input_activations": {
69 "num_bits": 4,
70 "type": "float",
71 "strategy": "tensor_group",
72 "group_size": 16,
73 "symmetric": True,
74 "dynamic": "local",
75 "observer": "minmax",
76 },
77 }
78 },
79 ),
80]
81
82# Save to disk in compressed-tensors format.
83SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
84
85# Apply quantization.
86oneshot(
87 model=model,
88 dataset=ds,
89 recipe=recipe,
90 max_seq_length=MAX_SEQUENCE_LENGTH,
91 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
92 output_dir=SAVE_DIR,
93)
94
95model.save_pretrained(SAVE_DIR, save_compressed=True)
96tokenizer.save_pretrained(SAVE_DIR)