Skip to main content

FP8 W8A8

vLLM 支持在 Nvidia H100 和 AMD MI300x 等 GPU 上利用硬件加速进行 FP8（8 位浮点数）的权重和激活量化。当前，仅 Hopper 和 Ada Lovelace GPU 对 W8A8 提供官方支持。对于 W8A16（仅权重为 FP8），可借助 Marlin 内核在 Ampere GPU 上得到支持。通过 FP8 对模型进行量化，能够使模型内存需求降低一半，吞吐量最多可提高 1.6 倍，同时对准确性的影响最小。

请访问 Hugging Face 上可与 vLLM 一起使用的流行 LLM 的量化 FP8 检查点集合。

硬件中通常支持的 FP8 类型有 2 种不同的表示形式，每种表示形式在不同的场景中都很有用：

  • E4M3：由 1 位符号位、4 位指数位和 3 位尾数组成。它可以存储高达 +/-448 和 nan 的值。

  • E5M2：由 1 位符号位、5 位指数位和 2 位尾数组成。它可以存储高达 +/-57344、+/- inf 和 nan 的值。增加动态范围的代价是存储值的精度较低。

注意：

FP8 计算支持计算能力大于 8.9 的 NVIDIA GPU (Ada Lovelace，Hopper)。FP8 模型将在计算能力大于 8.0 (Ampere) 的 GPU 上以仅权重为 W8A16 的形式运行，使用 FP8 Marlin。

使用在线动态量化快速入门​

使用 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
# INFO 06-10 17:55:42 model_runner.py:157] 加载模型权重占用了 0.1550 GB


result = model.generate("Hello, my name is")

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

安装​

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

pip install llmcompressor==0.1.0

量化过程​

量化过程涉及 3 个主要步骤:

  1. 加载模型

  2. 应用量化

  3. 评估 vLLM 的准确性

1. 加载模型​

使用封装了 AutoModelForCausalLM 的 SparseAutoModelForCausalLM 来保存和加载量化模型：

from llmcompressor.transformers import SparseAutoModelForCausalLM
from transformers import AutoTokenizer


MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"


model = SparseAutoModelForCausalLM.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
# 配置简单 PTQ 量化


recipe = QuantizationModifier(
  targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"])


# Apply the quantization algorithm.
# 应用量化算法。


oneshot(model=model, recipe=recipe)


# Save the model.
# 保存模型。


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.3

在 vllm 中加载并运行模型：

from vllm import LLM
model = LLM("./Meta-Llama-3-8B-Instruct-FP8-Dynamic")
model.generate("Hello my name is")

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

注意：

量化模型可能对 bos token 都存在很敏感。默认情况下，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。

已弃用的流程​

注意：

保留以下信息以供参考和搜索。下面所描述的量化方法已不再推荐使用，优先推荐使用上述的 llmcompressor 方法。

对于 FP8 的静态每张量离线量化，请安装 AutoFP8 库。

git clone https://github.com/neuralmagic/AutoFP8.git
pip install -e AutoFP8

该包引入了 AutoFP8ForCausalLM 和 BaseQuantizeConfig 对象，用于管理模型的压缩方式。

使用静态激活缩放因子进行离线量化​

您可以使用带有校准数据的 AutoFP8，通过启用 activation_scheme="static" 参数为权重和激活值生成每个张量的静态缩放。

from datasets import load_dataset
from transformers import AutoTokenizer
from auto_fp8 import AutoFP8ForCausalLM, BaseQuantizeConfig


pretrained_model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
quantized_model_dir = "Meta-Llama-3-8B-Instruct-FP8"


tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token


# Load and tokenize 512 dataset samples for calibration of activation scales
# 加载并标记 512 个数据集样本以校准激活量表


ds = load_dataset("mgoin/ultrachat_2k", split="train_sft").select(range(512))
examples = [tokenizer.apply_chat_template(batch["messages"], tokenize=False) for batch in ds]
examples = tokenizer(examples, padding=True, truncation=True, return_tensors="pt").to("cuda")


# Define quantization config with static activation scales
# 使用静态激活尺度定义量化配置


quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="static")


# Load the model, quantize, and save checkpoint
# 加载模型、量化并保存 checkpoint


model = AutoFP8ForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)
model.quantize(examples)
model.save_quantized(quantized_model_dir)

具有量化权重和激活值的模型 checkpoint 应当可以在 Meta-Llama-3-8B-Instruct-FP8/ 处获取。最后，您可以直接在 vLLM 中加载量化后的模型 checkpoint。

from vllm import LLM
model = LLM(model="Meta-Llama-3-8B-Instruct-FP8/")
# INFO 06-10 21:15:41 model_runner.py:159] Loading model weights took 8.4596 GB
# INFO 06-10 21:15:41 model_runner.py:159] 加载模型权重占用了 8.4596 GB


result = model.generate("Hello, my name is")