跳到內容

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