跳到内容

AMD Quark

量化可以有效地减少内存和带宽使用,加速计算并提高吞吐量,同时对精度损失极小。vLLM 可以利用 Quark,这是一个灵活且强大的量化工具包,来生成高性能的量化模型以在 AMD GPU 上运行。Quark 对 LLM 的权重、激活和 kv-cache 量化以及 AWQ、GPTQ、Rotation 和 SmoothQuant 等尖端量化算法提供了专门支持。

Quark 安装

在量化模型之前,您需要安装 Quark。可以使用 pip 安装最新发布的 Quark:

pip install amd-quark

有关更多安装详情,您可以参考 Quark 安装指南

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

pip install vllm git+https://github.com/EleutherAI/lm-evaluation-harness.git@206b7722158f58c35b7ffcd53b035fdbdda5126d#egg=lm-eval[api]

量化过程

安装 Quark 后,我们将使用一个示例来说明如何使用 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-cache 使用 FP8 每张量化,量化算法是 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. 量化模型并导出

然后我们可以应用量化。量化后,我们需要先冻结量化模型,然后再导出。请注意,我们需要以 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 离线量化的 MXFP4 和 MXFP6 模型,这些模型符合 Open Compute Project (OCP) 规范

该方案目前仅支持激活的动态量化。

示例用法,在安装最新 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 矩阵乘法执行的模拟,通过即时解量化 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