AMD Quark¶
量化可以在保持最小精度损失的同时,有效减少内存和带宽占用,加速计算并提高吞吐量。vLLM 可以利用 Quark(这一灵活且强大的量化工具包)来生成高性能的量化模型,从而在 AMD GPU 上运行。Quark 专门支持大语言模型量化,涵盖权重、激活值和 KV 缓存量化,以及 AWQ、GPTQ、Rotation 和 SmoothQuant 等前沿量化算法。
Quark 安装¶
在量化模型之前,需要先安装 Quark。可以通过 pip 安装最新版本的 Quark:
有关更多安装详情,请参考 Quark 安装指南。
此外,安装 vllm 和 lm-evaluation-harness 用于评估
量化过程¶
安装 Quark 后,我们将通过一个示例来说明如何使用它。Quark 量化流程可分为以下 5 个步骤:
- 加载模型
- 准备校准数据加载器
- 设置量化配置
- 量化模型并导出
- 在 vLLM 中评估
1. 加载模型¶
Quark 使用 Transformers 来获取模型和分词器。
代码
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_ID = "meta-llama/Llama-2-70b-chat-hf"
MAX_SEQ_LEN = 512
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map="auto",
dtype="auto",
)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, model_max_length=MAX_SEQ_LEN)
tokenizer.pad_token = tokenizer.eos_token
2. 准备校准数据加载器¶
Quark 使用 PyTorch Dataloader 加载校准数据。有关如何高效使用校准数据集的更多详情,请参考 添加校准数据集。
代码
from datasets import load_dataset
from torch.utils.data import DataLoader
BATCH_SIZE = 1
NUM_CALIBRATION_DATA = 512
# Load the dataset and get calibration data.
dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation")
text_data = dataset["text"][:NUM_CALIBRATION_DATA]
tokenized_outputs = tokenizer(
text_data,
return_tensors="pt",
padding=True,
truncation=True,
max_length=MAX_SEQ_LEN,
)
calib_dataloader = DataLoader(
tokenized_outputs['input_ids'],
batch_size=BATCH_SIZE,
drop_last=True,
)
3. 设置量化配置¶
我们需要设置量化配置。您可以查看 Quark 配置指南以获取详细信息。在此我们对权重、激活值和 KV 缓存使用 FP8 张量级(per-tensor)量化,使用的量化算法为 AutoSmoothQuant。
注意
请注意,量化算法需要一个 JSON 配置文件,该文件位于 Quark PyTorch 示例目录 examples/torch/language_modeling/llm_ptq/models 下。例如,Llama 的 AutoSmoothQuant 配置文件为 examples/torch/language_modeling/llm_ptq/models/llama/autosmoothquant_config.json。
代码
from quark.torch.quantization import (Config, QuantizationConfig,
FP8E4M3PerTensorSpec,
load_quant_algo_config_from_file)
# Define fp8/per-tensor/static spec.
FP8_PER_TENSOR_SPEC = FP8E4M3PerTensorSpec(
observer_method="min_max",
is_dynamic=False,
).to_quantization_spec()
# Define global quantization config, input tensors and weight apply FP8_PER_TENSOR_SPEC.
global_quant_config = QuantizationConfig(
input_tensors=FP8_PER_TENSOR_SPEC,
weight=FP8_PER_TENSOR_SPEC,
)
# Define quantization config for kv-cache layers, output tensors apply FP8_PER_TENSOR_SPEC.
KV_CACHE_SPEC = FP8_PER_TENSOR_SPEC
kv_cache_layer_names_for_llama = ["*k_proj", "*v_proj"]
kv_cache_quant_config = {
name: QuantizationConfig(
input_tensors=global_quant_config.input_tensors,
weight=global_quant_config.weight,
output_tensors=KV_CACHE_SPEC,
)
for name in kv_cache_layer_names_for_llama
}
layer_quant_config = kv_cache_quant_config.copy()
# Define algorithm config by config file.
LLAMA_AUTOSMOOTHQUANT_CONFIG_FILE = "examples/torch/language_modeling/llm_ptq/models/llama/autosmoothquant_config.json"
algo_config = load_quant_algo_config_from_file(LLAMA_AUTOSMOOTHQUANT_CONFIG_FILE)
EXCLUDE_LAYERS = ["lm_head"]
quant_config = Config(
global_quant_config=global_quant_config,
layer_quant_config=layer_quant_config,
kv_cache_quant_config=kv_cache_quant_config,
exclude=EXCLUDE_LAYERS,
algo_config=algo_config,
)
4. 量化模型并导出¶
接下来即可应用量化。量化完成后,在导出之前需要先对量化模型进行冻结(freeze)。请注意,导出模型时需采用 HuggingFace 的 safetensors 格式,更多关于导出格式的详情可参考 HuggingFace 格式导出。
代码
import torch
from quark.torch import ModelQuantizer, ModelExporter
from quark.torch.export import ExporterConfig, JsonExporterConfig
# Apply quantization.
quantizer = ModelQuantizer(quant_config)
quant_model = quantizer.quantize_model(model, calib_dataloader)
# Freeze quantized model to export.
freezed_model = quantizer.freeze(model)
# Define export config.
LLAMA_KV_CACHE_GROUP = ["*k_proj", "*v_proj"]
export_config = ExporterConfig(json_export_config=JsonExporterConfig())
export_config.json_export_config.kv_cache_group = LLAMA_KV_CACHE_GROUP
# Model: Llama-2-70b-chat-hf-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant
EXPORT_DIR = MODEL_ID.split("/")[1] + "-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant"
exporter = ModelExporter(config=export_config, export_dir=EXPORT_DIR)
with torch.no_grad():
exporter.export_safetensors_model(
freezed_model,
quant_config=quant_config,
tokenizer=tokenizer,
)
5. 在 vLLM 中评估¶
现在,您可以通过 LLM 入口直接加载并运行 Quark 量化后的模型。
代码
from vllm import LLM, SamplingParams
# Sample prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
# Create an LLM.
llm = LLM(
model="Llama-2-70b-chat-hf-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant",
kv_cache_dtype="fp8",
quantization="quark",
)
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
print("\nGenerated Outputs:\n" + "-" * 60)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}")
print(f"Output: {generated_text!r}")
print("-" * 60)
或者,您可以使用 lm_eval 来评估准确性。
lm_eval --model vllm \
--model_args pretrained=Llama-2-70b-chat-hf-w-fp8-a-fp8-kvcache-fp8-pertensor-autosmoothquant,kv_cache_dtype='fp8',quantization='quark' \
--tasks gsm8k
Quark 量化脚本¶
除了上述 Python API 示例外,Quark 还提供了一个 量化脚本,以便更方便地量化大语言模型。它支持使用各种不同的量化方案和优化算法来量化模型,能够即时导出量化模型并执行评估任务。使用该脚本,上述示例可以简化为:
python3 quantize_quark.py --model_dir meta-llama/Llama-2-70b-chat-hf \
--output_dir /path/to/output \
--quant_scheme w_fp8_a_fp8 \
--kv_cache_dtype fp8 \
--quant_algo autosmoothquant \
--num_calib_data 512 \
--model_export hf_format \
--tasks gsm8k
使用 OCP MX (MXFP4, MXFP6) 模型¶
vLLM 支持加载通过 AMD Quark 离线量化、符合 Open Compute Project (OCP) 规范的 MXFP4 和 MXFP6 模型。
该方案目前仅支持激活值的动态量化。
示例用法,在安装最新版 AMD Quark 后:
vllm serve fxmarty/qwen_1.5-moe-a2.7b-mxfp4 --tensor-parallel-size 1
# or, for a model using fp6 activations and fp4 weights:
vllm serve fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3 --tensor-parallel-size 1
对于不支持 OCP MX 原生操作的设备(例如 AMD Instinct MI325、MI300 和 MI250),可以运行 MXFP4/MXFP6 矩阵乘法执行的模拟,利用融合内核(fused kernel)即时将权重从 FP4/FP6 反量化为半精度。这对于使用 vLLM 评估 FP4/FP6 模型,或从约 2.5-4 倍的内存节省(与 float16 和 bfloat16 相比)中受益非常有用。
要生成使用 MXFP4 数据类型离线量化的模型,最简单的方法是使用 AMD Quark 的 量化脚本,示例如下:
python quantize_quark.py --model_dir Qwen/Qwen1.5-MoE-A2.7B-Chat \
--quant_scheme w_mxfp4_a_mxfp4 \
--output_dir qwen_1.5-moe-a2.7b-mxfp4 \
--skip_evaluation \
--model_export hf_format \
--group_size 32
当前的集成支持 用于权重或激活值的 FP4、FP6_E3M2、FP6_E2M3 的所有组合。
使用 Quark 量化逐层自动混合精度 (AMP) 模型¶
vLLM 还支持加载使用 AMD Quark 量化的逐层混合精度模型。目前支持 {MXFP4, FP8} 的混合方案,其中 FP8 指的是 FP8 张量级方案。计划在不久的将来支持更多混合精度方案,包括:
- 非量化线性层和/或 MoE 层作为每一层的选项,即 {MXFP4, FP8, BF16/FP16} 的混合。
- MXFP6 量化扩展,即 {MXFP4, MXFP6, FP8, BF16/FP16}。
尽管可以使用给定设备支持的最低精度(例如 AMD Instinct MI355 的 MXFP4,AMD Instinct MI300 的 FP8)来最大化服务吞吐量,但这些激进的方案可能会损害目标任务在量化后的精度恢复。混合精度可以在最大化精度和吞吐量之间取得平衡。
生成和部署使用 AMD Quark 量化的混合精度模型分为两个步骤,如下所示:
1. 在 AMD Quark 中使用混合精度量化模型¶
首先,搜索给定 LLM 模型的逐层混合精度配置,然后使用 AMD Quark 进行量化。稍后我们将提供使用 Quark API 的详细教程。
作为示例,我们提供了一些现成的量化混合精度模型,以展示其在 vLLM 中的用法及精度收益。它们是:
- amd/Llama-2-70b-chat-hf-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
- amd/Mixtral-8x7B-Instruct-v0.1-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
- amd/Qwen3-8B-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
2. 在 vLLM 中推理量化后的混合精度模型¶
使用 AMD Quark 混合精度量化的模型可以在 vLLM 中直接加载,并可通过 lm-evaluation-harness 进行如下评估: