跳到内容

AMD Quark

量化可以在保持最小精度损失的同时,有效减少内存和带宽占用,加速计算并提高吞吐量。vLLM 可以利用 Quark(这一灵活且强大的量化工具包)来生成高性能的量化模型,从而在 AMD GPU 上运行。Quark 专门支持大语言模型量化,涵盖权重、激活值和 KV 缓存量化,以及 AWQ、GPTQ、Rotation 和 SmoothQuant 等前沿量化算法。

Quark 安装

在量化模型之前,需要先安装 Quark。可以通过 pip 安装最新版本的 Quark:

pip install amd-quark

有关更多安装详情,请参考 Quark 安装指南

此外,安装 vllmlm-evaluation-harness 用于评估

pip install vllm "lm-eval[api]>=0.4.12"

量化过程

安装 Quark 后,我们将通过一个示例来说明如何使用它。Quark 量化流程可分为以下 5 个步骤:

  1. 加载模型
  2. 准备校准数据加载器
  3. 设置量化配置
  4. 量化模型并导出
  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 进行如下评估:

lm_eval --model vllm \
    --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 \
    --tasks mmlu \
    --batch_size auto