跳到内容

FP8 W8A8

vLLM支持FP8(8位浮点)权重和激活量化,利用Nvidia H100和AMD MI300x等GPU的硬件加速。目前,W8A8只正式支持Hopper和Ada Lovelace GPU。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模型将在计算能力 > 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层进行量化,该方案采用

  • 权重的静态、逐通道量化
  • 激活的动态、逐token量化

由于简单的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. 评估精度

安装vllmlm-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仓库中提交问题。

在线动态量化

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

在此模式下,所有Linear模块(最终的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位之前会加载原始精度模型,因此您需要足够的内存来加载整个模型。