@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 395 lines 17 KB
Rendered Raw
1 # AMD Quark
2
3 Quantization can effectively reduce memory and bandwidth usage, accelerate computation and improve
4 throughput while with minimal accuracy loss. vLLM can leverage [Quark](https://quark.docs.amd.com/latest/),
5 the flexible and powerful quantization toolkit, to produce performant quantized models to run on AMD GPUs. Quark has specialized support for quantizing large language models with weight,
6 activation and kv-cache quantization and cutting-edge quantization algorithms like
7 AWQ, GPTQ, Rotation and SmoothQuant.
8
9 ## Quark Installation
10
11 Before quantizing models, you need to install Quark. The latest release of Quark can be installed with pip:
12
13 ```bash
14 pip install amd-quark
15 ```
16
17 You can refer to [Quark installation guide](https://quark.docs.amd.com/latest/install.html)
18 for more installation details.
19
20 Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
21
22 ```bash
23 pip install vllm "lm-eval[api]>=0.4.12"
24 ```
25
26 ## Quantization Process
27
28 After installing Quark, we will use an example to illustrate how to use Quark.
29 The Quark quantization process can be listed for 5 steps as below:
30
31 1. Load the model
32 2. Prepare the calibration dataloader
33 3. Set the quantization configuration
34 4. Quantize the model and export
35 5. Evaluation in vLLM
36
37 ### 1. Load the Model
38
39 Quark uses [Transformers](https://huggingface.co/docs/transformers/en/index)
40 to fetch model and tokenizer.
41
42 ??? code
43
44 ```python
45 from transformers import AutoTokenizer, AutoModelForCausalLM
46
47 MODEL_ID = "meta-llama/Llama-2-70b-chat-hf"
48 MAX_SEQ_LEN = 512
49
50 model = AutoModelForCausalLM.from_pretrained(
51 MODEL_ID,
52 device_map="auto",
53 dtype="auto",
54 )
55 model.eval()
56
57 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, model_max_length=MAX_SEQ_LEN)
58 tokenizer.pad_token = tokenizer.eos_token
59 ```
60
61 ### 2. Prepare the Calibration Dataloader
62
63 Quark uses the [PyTorch Dataloader](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)
64 to load calibration data. For more details about how to use calibration datasets efficiently, please refer
65 to [Adding Calibration Datasets](https://quark.docs.amd.com/latest/pytorch/calibration_datasets.html).
66
67 ??? code
68
69 ```python
70 from datasets import load_dataset
71 from torch.utils.data import DataLoader
72
73 BATCH_SIZE = 1
74 NUM_CALIBRATION_DATA = 512
75
76 # Load the dataset and get calibration data.
77 dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation")
78 text_data = dataset["text"][:NUM_CALIBRATION_DATA]
79
80 tokenized_outputs = tokenizer(
81 text_data,
82 return_tensors="pt",
83 padding=True,
84 truncation=True,
85 max_length=MAX_SEQ_LEN,
86 )
87 calib_dataloader = DataLoader(
88 tokenized_outputs['input_ids'],
89 batch_size=BATCH_SIZE,
90 drop_last=True,
91 )
92 ```
93
94 ### 3. Set the Quantization Configuration
95
96 We need to set the quantization configuration, you can check
97 [quark config guide](https://quark.docs.amd.com/latest/pytorch/user_guide_config_description.html)
98 for further details. Here we use FP8 per-tensor quantization on weight, activation,
99 kv-cache and the quantization algorithm is AutoSmoothQuant.
100
101 !!! note
102 Note the quantization algorithm needs a JSON config file and the config file is located in
103 [Quark PyTorch examples](https://quark.docs.amd.com/latest/pytorch/pytorch_examples.html),
104 under the directory `examples/torch/language_modeling/llm_ptq/models`. For example,
105 AutoSmoothQuant config file for Llama is
106 `examples/torch/language_modeling/llm_ptq/models/llama/autosmoothquant_config.json`.
107
108 ??? code
109
110 ```python
111 from quark.torch.quantization import (Config, QuantizationConfig,
112 FP8E4M3PerTensorSpec,
113 load_quant_algo_config_from_file)
114
115 # Define fp8/per-tensor/static spec.
116 FP8_PER_TENSOR_SPEC = FP8E4M3PerTensorSpec(
117 observer_method="min_max",
118 is_dynamic=False,
119 ).to_quantization_spec()
120
121 # Define global quantization config, input tensors and weight apply FP8_PER_TENSOR_SPEC.
122 global_quant_config = QuantizationConfig(
123 input_tensors=FP8_PER_TENSOR_SPEC,
124 weight=FP8_PER_TENSOR_SPEC,
125 )
126
127 # Define quantization config for kv-cache layers, output tensors apply FP8_PER_TENSOR_SPEC.
128 KV_CACHE_SPEC = FP8_PER_TENSOR_SPEC
129 kv_cache_layer_names_for_llama = ["*k_proj", "*v_proj"]
130 kv_cache_quant_config = {
131 name: QuantizationConfig(
132 input_tensors=global_quant_config.input_tensors,
133 weight=global_quant_config.weight,
134 output_tensors=KV_CACHE_SPEC,
135 )
136 for name in kv_cache_layer_names_for_llama
137 }
138 layer_quant_config = kv_cache_quant_config.copy()
139
140 # Define algorithm config by config file.
141 LLAMA_AUTOSMOOTHQUANT_CONFIG_FILE = "examples/torch/language_modeling/llm_ptq/models/llama/autosmoothquant_config.json"
142 algo_config = load_quant_algo_config_from_file(LLAMA_AUTOSMOOTHQUANT_CONFIG_FILE)
143
144 EXCLUDE_LAYERS = ["lm_head"]
145 quant_config = Config(
146 global_quant_config=global_quant_config,
147 layer_quant_config=layer_quant_config,
148 kv_cache_quant_config=kv_cache_quant_config,
149 exclude=EXCLUDE_LAYERS,
150 algo_config=algo_config,
151 )
152 ```
153
154 ### 4. Quantize the Model and Export
155
156 Then we can apply the quantization. After quantizing, we need to freeze the
157 quantized model first before exporting. Note that we need to export model with format of
158 HuggingFace `safetensors`, you can refer to
159 [HuggingFace format exporting](https://quark.docs.amd.com/latest/pytorch/export/quark_export_hf.html)
160 for more exporting format details.
161
162 ??? code
163
164 ```python
165 import torch
166 from quark.torch import ModelQuantizer, ModelExporter
167 from quark.torch.export import ExporterConfig, JsonExporterConfig
168
169 # Apply quantization.
170 quantizer = ModelQuantizer(quant_config)
171 quant_model = quantizer.quantize_model(model, calib_dataloader)
172
173 # Freeze quantized model to export.
174 freezed_model = quantizer.freeze(model)
175
176 # Define export config.
177 LLAMA_KV_CACHE_GROUP = ["*k_proj", "*v_proj"]
178 export_config = ExporterConfig(json_export_config=JsonExporterConfig())
179 export_config.json_export_config.kv_cache_group = LLAMA_KV_CACHE_GROUP
180
181 # Model: Llama-2-70b-chat-hf-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant
182 EXPORT_DIR = MODEL_ID.split("/")[1] + "-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant"
183 exporter = ModelExporter(config=export_config, export_dir=EXPORT_DIR)
184 with torch.no_grad():
185 exporter.export_safetensors_model(
186 freezed_model,
187 quant_config=quant_config,
188 tokenizer=tokenizer,
189 )
190 ```
191
192 ### 5. Evaluation in vLLM
193
194 Now, you can load and run the Quark quantized model directly through the LLM entrypoint:
195
196 ??? code
197
198 ```python
199 from vllm import LLM, SamplingParams
200
201 # Sample prompts.
202 prompts = [
203 "Hello, my name is",
204 "The president of the United States is",
205 "The capital of France is",
206 "The future of AI is",
207 ]
208 # Create a sampling params object.
209 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
210
211 # Create an LLM.
212 llm = LLM(
213 model="Llama-2-70b-chat-hf-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant",
214 kv_cache_dtype="fp8",
215 quantization="quark",
216 )
217 # Generate texts from the prompts. The output is a list of RequestOutput objects
218 # that contain the prompt, generated text, and other information.
219 outputs = llm.generate(prompts, sampling_params)
220 # Print the outputs.
221 print("\nGenerated Outputs:\n" + "-" * 60)
222 for output in outputs:
223 prompt = output.prompt
224 generated_text = output.outputs[0].text
225 print(f"Prompt: {prompt!r}")
226 print(f"Output: {generated_text!r}")
227 print("-" * 60)
228 ```
229
230 Or, you can use `lm_eval` to evaluate accuracy:
231
232 ```bash
233 lm_eval --model vllm \
234 --model_args pretrained=Llama-2-70b-chat-hf-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant,kv_cache_dtype='fp8',quantization='quark' \
235 --tasks gsm8k
236 ```
237
238 ## Quark Quantization Script
239
240 In addition to the example of Python API above, Quark also offers a
241 [quantization script](https://quark.docs.amd.com/latest/pytorch/example_quark_torch_llm_ptq.html)
242 to quantize large language models more conveniently. It supports quantizing models with variety
243 of different quantization schemes and optimization algorithms. It can export the quantized model
244 and run evaluation tasks on the fly. With the script, the example above can be:
245
246 ```bash
247 python3 quantize_quark.py --model_dir meta-llama/Llama-2-70b-chat-hf \
248 --output_dir /path/to/output \
249 --quant_scheme w_fp8_a_fp8 \
250 --kv_cache_dtype fp8 \
251 --quant_algo autosmoothquant \
252 --num_calib_data 512 \
253 --model_export hf_format \
254 --tasks gsm8k
255 ```
256
257 ## Using OCP MX (MXFP4, MXFP6) models
258
259 vLLM supports loading MXFP4 and MXFP6 models quantized offline through AMD Quark, compliant with [Open Compute Project (OCP) specification](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
260
261 The scheme currently only supports dynamic quantization for activations.
262
263 Example usage, after installing the latest AMD Quark release:
264
265 ```bash
266 vllm serve fxmarty/qwen_1.5-moe-a2.7b-mxfp4 --tensor-parallel-size 1
267 # or, for a model using fp6 activations and fp4 weights:
268 vllm serve fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3 --tensor-parallel-size 1
269 ```
270
271 A simulation of the matrix multiplication execution in MXFP4/MXFP6 can be run on devices that do not support OCP MX operations natively (e.g. AMD Instinct MI325, MI300 and MI250), dequantizing weights from FP4/FP6 to half precision on the fly, using a fused kernel. This is useful e.g. to evaluate FP4/FP6 models using vLLM, or alternatively to benefit from the ~2.5-4x memory savings (compared to float16 and bfloat16).
272
273 To generate offline models quantized using MXFP4 data type, the easiest approach is to use AMD Quark's [quantization script](https://quark.docs.amd.com/latest/pytorch/example_quark_torch_llm_ptq.html), as an example:
274
275 ```bash
276 python quantize_quark.py --model_dir Qwen/Qwen1.5-MoE-A2.7B-Chat \
277 --quant_scheme w_mxfp4_a_mxfp4 \
278 --output_dir qwen_1.5-moe-a2.7b-mxfp4 \
279 --skip_evaluation \
280 --model_export hf_format \
281 --group_size 32
282 ```
283
284 The current integration supports [all combination of FP4, FP6_E3M2, FP6_E2M3](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py) used for either weights or activations.
285
286 ## Using Quark Quantized layerwise Auto Mixed Precision (AMP) Models
287
288 vLLM also supports loading layerwise mixed precision model quantized using AMD Quark. Currently, mixed scheme of {MXFP4, FP8} is supported, where FP8 here denotes for FP8 per-tensor scheme. More mixed precision schemes are planned to be supported in a near future, including
289
290 - Unquantized Linear and/or MoE layer(s) as an option for each layer, i.e., mixed of {MXFP4, FP8, BF16/FP16}
291 - MXFP6 quantization extension, i.e., {MXFP4, MXFP6, FP8, BF16/FP16}
292
293 Although one can maximize serving throughput using the lowest precision supported on a given device (e.g. MXFP4 for AMD Instinct MI355, FP8 for AMD Instinct MI300), these aggressive schemes can be detrimental to accuracy recovering from quantization on target tasks. Mixed precision allows to strike a balance between maximizing accuracy and throughput.
294
295 There are two steps to generate and deploy a mixed precision model quantized with AMD Quark, as shown below.
296
297 ### 1. Quantize a model using mixed precision in AMD Quark
298
299 Firstly, the layerwise mixed-precision configuration for a given LLM model is searched and then quantized using AMD Quark. We will provide a detailed tutorial with Quark APIs later.
300
301 As examples, we provide some ready-to-use quantized mixed precision model to show the usage in vLLM and the accuracy benefits. They are:
302
303 - amd/Llama-2-70b-chat-hf-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
304 - amd/Mixtral-8x7B-Instruct-v0.1-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
305 - amd/Qwen3-8B-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
306
307 ### 2. inference the quantized mixed precision model in vLLM
308
309 Models quantized with AMD Quark using mixed precision can natively be reload in vLLM, and e.g. evaluated using lm-evaluation-harness as follows:
310
311 ```bash
312 lm_eval --model vllm \
313 --model_args pretrained=amd/Llama-2-70b-chat-hf-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8,tensor_parallel_size=4,dtype=auto,gpu_memory_utilization=0.8,trust_remote_code=False \
314 --tasks mmlu \
315 --batch_size auto
316 ```
317
318 ## Online Quantization
319
320 All the workflows above are *offline* quantization: you run a script, write a new quantized checkpoint to disk, and later load it for serving. This produces the most accurate and smallest checkpoints because it can run accuracy-recovery algorithms (rotation, SmoothQuant, GPTQ, AWQ) during export, but it requires a separate quantization step and a second copy of the model on disk.
321
322 *Online* quantization instead quantizes the weights at load time, directly from a high-precision checkpoint, and offers several advantages over the offline flow:
323
324 - **No export step** — serve directly from the original `bf16`/`fp16` checkpoint; no separate quantization run before deployment.
325 - **No extra disk footprint** — nothing new is written to disk, so there is no second copy of the model to store or manage.
326 - **No calibration data** — activations are scaled dynamically at runtime, so no calibration dataset is needed.
327 - **Fast iteration** — switch schemes or per-layer selections instantly by changing a config, without re-exporting.
328
329 ### vLLM online quantization
330
331 vLLM has [built-in online quantization](online.md) that skips the export step: it loads an ordinary high-precision (`bf16`) checkpoint and quantizes each layer's weights to schemes such as FP8 or MXFP4 at load time, without a pre-quantized checkpoint or calibration data. No new checkpoint is written to disk. See [Online Quantization](online.md) for the supported schemes and configuration.
332
333 ### Quark online quantization
334
335 AMD Quark provides its own online path, exposed as the `quark_online` quantization backend, for users who want parity with Quark's offline export and Quark's per-layer mixed-precision configs. Like vLLM's built-in online quantization, it quantizes weights inside the weight-loading hook just before serving and writes no new checkpoint.
336
337 Its quant math is aligned byte-for-byte with Quark's offline export, so what you validate online is what you get offline. Use it to *find* the scheme and per-layer selection you want.
338
339 Compared to vLLM's built-in online quantization, the `quark_online` plugin adds:
340
341 - **Flexible config parsing** — a terse config expands into Quark's verbose per-layer config, delegating all per-layer matching to a real `QuarkConfig`.
342 - **Per-layer / mixed schemes** — dispatch a different method per layer (e.g. MXFP4 experts with FP8 attention on an MoE model). Each online method subclasses the matching offline scheme, so a load-time quantized layer runs the identical inference kernel as an offline one.
343 - **Re-quantizing an already-quantized checkpoint** — an FP8 block-scale checkpoint (e.g. DeepSeek-R1) is dequantized and re-quantized to a target scheme layer-locally at load time, with no new checkpoint.
344
345 The plugin ships in AMD Quark — no fork of vLLM, no patched checkpoint format; see the [Quark documentation](https://quark.docs.amd.com/latest/) for details. Three presets ship ready to use:
346
347 | Preset key | Scheme |
348 | --- | --- |
349 | `ptpc_fp8` | FP8 E4M3, per-channel weight + dynamic per-token activation |
350 | `mxfp4` | MXFP4, per-group (group size 32) with E8M0 block scale |
351 | `linear_ptpc_fp8_moe_mxfp4` | Mixed: attention in FP8, MoE experts in MXFP4 |
352
353 ### Python API
354
355 ```python
356 from vllm import LLM, SamplingParams
357 from quark.online_quantization.vllm import HF_QUANTIZATION_CONFIGS
358
359 llm = LLM(
360 model="Qwen/Qwen3-30B-A3B-Thinking-2507",
361 quantization="quark_online",
362 hf_overrides=HF_QUANTIZATION_CONFIGS["ptpc_fp8"],
363 enforce_eager=True,
364 tensor_parallel_size=1,
365 )
366 out = llm.generate(["The capital of France is"], SamplingParams(temperature=0.0, max_tokens=100))
367 print(out)
368 ```
369
370 ### Native vLLM CLI
371
372 ```bash
373 export VLLM_PLUGINS="${VLLM_PLUGINS:-quark_online_quant}"
374 ONLINE_QUANT_CONFIG='{"online_quant_config": {"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head"]}}'
375
376 vllm serve Qwen/Qwen3-8B \
377 --trust-remote-code \
378 --tensor-parallel-size 1 \
379 --additional-config "$ONLINE_QUANT_CONFIG"
380 ```
381
382 ### Re-quantizing an offline checkpoint
383
384 No extra arguments — the same `hf_overrides` detects the checkpoint's existing config and merges it automatically:
385
386 ```python
387 llm = LLM(
388 model="deepseek-ai/DeepSeek-R1", # ships quant_method: "fp8"
389 quantization="quark_online",
390 hf_overrides=HF_QUANTIZATION_CONFIGS["ptpc_fp8"],
391 tensor_parallel_size=8,
392 )
393 ```
394
395 The plugin can be disabled with `QUARK_DISABLE_VLLM_PLUGIN=1`, and it steps aside automatically if another platform plugin (e.g. ATOM) has already taken over the backend.