跳至内容

FP8 W8A8

vLLM支持在Nvidia H100和AMD MI300x等GPU上使用硬件加速进行FP8（8位浮点）权重和激活量化。 目前，仅Hopper和Ada Lovelace架构GPU官方支持W8A8量化。 对于Ampere架构GPU，可通过Marlin内核支持W8A16（仅权重FP8）量化。 使用FP8量化模型可将内存需求降低2倍，并在精度影响最小的情况下实现高达1.6倍的吞吐量提升。

请访问HF上的已量化FP8热门LLM检查点合集，可直接与vLLM配合使用。

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

  • E4M3: 包含1个符号位、4个指数位和3个尾数位。可以存储高达+/-448的值和nan。
  • E5M2: 由1个符号位、5个指数位和2个尾数位组成。可存储高达+/-57344的值、+/-inf和nan。增加动态范围的代价是存储值的精度降低。

注意

FP8计算支持在计算能力大于8.9（Ada Lovelace、Hopper）的NVIDIA GPU上运行。 FP8模型将在计算能力大于8.0（Ampere）的GPU上以仅权重W8A16模式运行，利用FP8 Marlin技术。

安装

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

pip install llmcompressor

量化处理过程

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

  1. 加载模型
  2. 应用量化技术
  3. 在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", torch_dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

2. 应用量化

对于FP8量化，我们可以通过简单的RTN量化恢复精度。我们建议使用FP8_DYNAMIC方案针对所有Linear层，该方案使用：

  • 对权重进行静态、逐通道量化
  • 对激活值进行动态的逐令牌量化

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

from llmcompressor.transformers 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==0.4.4

在vllm中加载并运行模型：

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

使用lm_eval评估准确率（例如在gsm8k的250个样本上）：

注意

量化模型可能对bos标记的存在很敏感。lm_eval默认不会添加bos标记，因此在运行评估时请确保包含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。

在线动态量化

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

在此模式下，所有线性模块（除最终的lm_head外）的权重都会被量化为FP8_E4M3精度，并采用逐张量缩放。在前向传播过程中会动态计算激活值的最小最大值，以提供高精度的动态逐张量缩放。因此，该模式下的延迟改善效果有限。

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

警告

目前，我们会在将模型量化为8位之前以原始精度加载模型，因此您需要有足够的内存来加载整个模型。