跳到内容

FP8 W8A8

vLLM 支持使用 FP8(8 位浮点数)的权重和激活量化,并在 Nvidia H100 和 AMD MI300x 等 GPU 上进行硬件加速。目前,W8A8 仅官方支持 Hopper 和 Ada Lovelace GPU。Ampere GPU 通过利用 Marlin 内核支持 W8A16(仅权重 FP8)。FP8 量化模型可将模型内存需求减少 2 倍,并将吞吐量提高高达 1.6 倍,同时对准确性的影响最小。

请访问 Hugging Face 上精选的 vLLM 可用的量化 FP8 流行 LLM 检查点集合

硬件通常支持的 FP8 类型有两种不同的表示形式,它们各自适用于不同的场景:

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

注意

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",
    dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

2. 应用量化

对于 FP8 量化,我们可以通过简单的 RTN 量化来恢复准确性。我们建议使用 FP8_DYNAMIC 方案,该方案针对所有 Linear 层,它使用:

  • 对权重的静态、每通道量化
  • 对激活的动态、每令牌量化

由于简单的 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. 评估准确性

安装 vllmlm-evaluation-harness 以进行评估。

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

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 评估准确性(例如,在 gsm8k 的 250 个样本上)。

注意

量化模型可能对 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" 来启用此功能。

在此模式下,所有 Linear 模块(lm_head 除外)的权重将被量化到 FP8_E4M3 精度,并使用每张量(per-tensor)缩放因子。激活值在每次前向传递期间计算其最小值和最大值,以提供动态的每张量缩放因子以获得高精度。因此,在此模式下,延迟提升有限。

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)

警告

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