FP8 W8A8¶

vLLM 支持在 Nvidia H100 和 AMD MI300x 等 GPU 上利用硬件加速进行 FP8（8 位浮点数）权重和激活量化。目前，仅 Hopper 和 Ada Lovelace 架构的 GPU 官方支持 W8A8。Turing/Ampere 架构的 GPU 可通过 Marlin 内核支持 W8A16（仅权重 FP8）。模型进行 FP8 量化后，内存需求可减少 2 倍，吞吐量提升最高可达 1.6 倍，且对准确性的影响极小。

请访问 HF 合集，获取适用于 vLLM 的主流 LLM 量化 FP8 检查点。

硬件通常支持的 FP8 类型有两种不同的表示形式，每种在不同场景下各有用处：

E4M3：由 1 位符号位、4 位指数位和 3 位尾数位组成。它可以存储最大至 +/-448 的数值以及 nan。
E5M2：由 1 位符号位、5 位指数位和 2 位尾数位组成。它可以存储最大至 +/-57344 的数值、+/- inf 以及 nan。这种动态范围增加的代价是存储数值的精度较低。

注意

计算能力 >= 8.9（Ada Lovelace, Hopper）的 NVIDIA GPU 支持 FP8 计算。FP8 模型可在计算能力 >= 7.5（Turing）的 GPU 上以仅权重 W8A16 模式运行，利用 FP8 Marlin 内核。

安装¶

要使用 vLLM 生成高性能的 FP8 量化模型，需要安装 llm-compressor 库。

pip install llmcompressor

量化过程¶

量化过程包含三个主要步骤：

加载模型
应用量化
在 vLLM 中评估精度

1. 加载模型¶

使用标准的 transformers AutoModel 类加载您的模型和分词器

from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

2. 应用量化¶

对于 FP8 量化，我们可以通过简单的 RTN（舍入到最近值）量化来恢复准确性。我们建议对所有 Linear 层使用 FP8_DYNAMIC 方案，该方案使用：

对权重进行静态、按通道（per-channel）量化
对激活值进行动态、按 Token（per-token）量化

由于简单的 RTN 不需要数据进行权重量化，且激活值是动态量化的，因此此量化流程无需任何校准数据。

代码

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

# Configure the simple PTQ quantization
recipe = QuantizationModifier(
    targets="Linear",
    scheme="FP8_DYNAMIC",
    ignore=["lm_head"],
)

# Apply the quantization algorithm.
oneshot(model=model, recipe=recipe)

# Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)

3. 评估准确性¶

安装 vllm 和 lm-evaluation-harness 进行评估。

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

在 vllm 中加载并运行模型。

from vllm import LLM

llm = LLM("./Meta-Llama-3-8B-Instruct-FP8-Dynamic")
result = llm.generate("Hello my name is")
print(result[0].outputs[0].text)

使用 lm_eval 评估准确性（例如在 250 个 gsm8k 样本上）。

注意

量化模型可能对 bos token 的存在较为敏感。lm_eval 默认不会添加 bos token，因此请确保在运行评估时包含 add_bos_token=True 参数。

MODEL=$PWD/Meta-Llama-3-8B-Instruct-FP8-Dynamic
lm_eval \
  --model vllm \
  --model_args pretrained=$MODEL,add_bos_token=True \
  --tasks gsm8k  --num_fewshot 5 --batch_size auto --limit 250

以下是结果得分的示例：

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
| --- |------:| -------------- |-----:| --------- | - |----:| - |-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.768|±  |0.0268|
|     |       |strict-match    |     5|exact_match|↑  |0.768|±  |0.0268|

故障排除和支持¶

如果您遇到任何问题或有功能需求，请在 vllm-project/llm-compressor GitHub 仓库中提交 issue。

在线动态量化¶

无需任何校准数据，即可通过 vLLM 将原始精度的 BF16/FP16 模型动态量化为 FP8。您可以通过在命令行中指定 --quantization="fp8" 或在 LLM 构造函数中设置 quantization="fp8" 来启用此功能。

在这种模式下，所有线性模块（最终的 lm_head 除外）的权重都会被量化为具有张量级缩放（per-tensor scale）的 FP8_E4M3 精度。激活值的最大值和最小值会在每次前向传播过程中计算，以提供动态的张量级缩放，从而确保高精度。因此，此模式下的延迟改善有限。

from vllm import LLM

llm = LLM("facebook/opt-125m", quantization="fp8")
# INFO 06-10 17:55:42 model_runner.py:157] Loading model weights took 0.1550 GB
result = llm.generate("Hello, my name is")
print(result[0].outputs[0].text)