@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 217 lines 6.42 KB
Rendered Raw
1 # INT8 W4A8
2
3 vLLM supports quantizing weights to INT4 and activations to INT8 for memory savings and inference acceleration.
4 This quantization method is particularly useful for reducing model size while maintaining good performance.
5
6 ## Prerequisites
7
8 To use INT8 W4A8 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library.
9
10 ```bash
11 (venv-llm-compressor) pip install llmcompressor
12 ```
13
14 Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
15
16 ```bash
17 (venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
18 ```
19
20 Please use separate environments for vLLM and llm-compressor as they might not work together.
21
22 ## Quantization Process
23
24 The quantization process involves four main steps:
25
26 1. Loading the model
27 2. Preparing calibration data
28 3. Applying quantization
29 4. Evaluating accuracy in vLLM
30
31 ### 1. Loading the Model
32
33 Load your model and tokenizer using the standard `transformers` AutoModel classes:
34
35 ```python
36 from transformers import AutoTokenizer, AutoModelForCausalLM
37
38 MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
39 model = AutoModelForCausalLM.from_pretrained(
40 MODEL_ID,
41 dtype="auto",
42 )
43 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
44 ```
45
46 ### 2. Preparing Calibration Data
47
48 When quantizing activations to INT8 and weights to INT4, you need sample data to estimate the activation scales.
49 It's best to use calibration data that closely matches your deployment data.
50 For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
51
52 ```python
53 from datasets import load_dataset
54
55 NUM_CALIBRATION_SAMPLES = 512
56 MAX_SEQUENCE_LENGTH = 2048
57
58 # Load and preprocess the dataset
59 ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
60 ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
61
62 def preprocess(example):
63 return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
64 ds = ds.map(preprocess)
65
66 def tokenize(sample):
67 return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
68 ds = ds.map(tokenize, remove_columns=ds.column_names)
69 ```
70
71 ### 3. Applying Quantization
72
73 Now, apply the quantization algorithms.
74
75 The following recipes create W4A8 models (int4 weights, int8 activations). On Arm® CPUs, this is accelerated through [KleidiAI](https://github.com/ARM-software/kleidiai).
76
77 Use groupwise for best accuracy, and channelwise for best inference performance.
78
79 === "Groupwise"
80
81 ```python
82 from llmcompressor import oneshot
83 from llmcompressor.modifiers.quantization import GPTQModifier
84
85 # Configure the quantization algorithms
86 recipe = [
87 GPTQModifier(
88 targets="Linear",
89 scheme="W4A8",
90 ignore=["lm_head"],
91 dampening_frac=0.01
92 ),
93 ]
94
95 # Apply quantization
96 oneshot(
97 model=model,
98 dataset=ds,
99 recipe=recipe,
100 max_seq_length=MAX_SEQUENCE_LENGTH,
101 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
102 )
103
104 # Save the compressed model: Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token
105 SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A8-G128-Dynamic-Per-Token"
106 model.save_pretrained(SAVE_DIR, save_compressed=True)
107 tokenizer.save_pretrained(SAVE_DIR)
108 ```
109
110 === "Channelwise"
111
112 ```python
113 from llmcompressor import oneshot
114 from llmcompressor.modifiers.quantization import GPTQModifier
115 from compressed_tensors.quantization import QuantizationStrategy, QuantizationType
116
117 scheme = {
118 "targets": ["Linear"],
119 "weights": {
120 "num_bits": 4,
121 "type": QuantizationType.INT,
122 "strategy": QuantizationStrategy.CHANNEL,
123 "symmetric": True,
124 "dynamic": False,
125 "group_size": None,
126 },
127 "input_activations": {
128 "num_bits": 8,
129 "type": QuantizationType.INT,
130 "strategy": QuantizationStrategy.TOKEN,
131 "dynamic": True,
132 "symmetric": False,
133 "observer": None,
134 },
135 "output_activations": None,
136 }
137
138 recipe = [
139 GPTQModifier(
140 targets="Linear",
141 config_groups={"group_0": scheme},
142 ignore=["lm_head"],
143 dampening_frac=0.01,
144 ),
145 ]
146
147 oneshot(
148 model=model,
149 dataset=ds,
150 recipe=recipe,
151 max_seq_length=MAX_SEQUENCE_LENGTH,
152 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
153 )
154
155 # Save the compressed model: Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token
156 SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A8-Channelwise-Dynamic-Per-Token"
157 model.save_pretrained(SAVE_DIR, save_compressed=True)
158 tokenizer.save_pretrained(SAVE_DIR)
159 ```
160
161 ### 4. Evaluating Accuracy
162
163 === "Groupwise"
164
165 After quantization, you can load and run the model in vLLM:
166
167 ```python
168 from vllm import LLM
169
170 llm = LLM("./Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token")
171 ```
172
173 To evaluate accuracy, you can use `lm_eval`:
174
175 ```bash
176 lm_eval --model vllm \
177 --model_args pretrained="./Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token",add_bos_token=true \
178 --tasks gsm8k \
179 --num_fewshot 5 \
180 --limit 250 \
181 --batch_size 'auto'
182 ```
183
184 === "Channelwise"
185
186 After quantization, you can load and run the model in vLLM:
187
188 ```python
189 from vllm import LLM
190
191 llm = LLM("./Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token")
192 ```
193
194 To evaluate accuracy, you can use `lm_eval`:
195
196 ```bash
197 lm_eval --model vllm \
198 --model_args pretrained="./Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token",add_bos_token=true \
199 --tasks gsm8k \
200 --num_fewshot 5 \
201 --limit 250 \
202 --batch_size 'auto'
203 ```
204
205 !!! note
206 Quantized models can be sensitive to the presence of the `bos` token. Make sure to include the `add_bos_token=True` argument when running evaluations.
207
208 ## Best Practices
209
210 - Start with 512 samples for calibration data (increase if accuracy drops)
211 - Use a sequence length of 2048 as a starting point
212 - Employ the chat template or instruction template that the model was trained with
213 - If you've fine-tuned a model, consider using a sample of your training data for calibration
214
215 ## Troubleshooting and Support
216
217 If you encounter any issues or have feature requests, please open an issue on the [vllm-project/llm-compressor](https://github.com/vllm-project/llm-compressor/issues) GitHub repository.