跳到内容

fp8 权重和激活量化

llmcompressor 支持将权重和激活量化到 fp8,以节省内存并使用 vllm 加速推理。

fp8 计算支持 Nvidia GPU,其计算能力 > 8.9 (Ada Lovelace, Hopper)。

安装

首先,安装

pip install llmcompressor

快速入门

该示例包含一个用于应用量化算法的端到端脚本。

python3 llama3_example.py

生成的模型 Meta-Llama-3-8B-Instruct-FP8-Dynamic 已准备好加载到 vLLM 中。

代码演练

现在,我们将逐步介绍示例中的代码。共有三个步骤:1) 加载模型 2) 应用量化 3) 在 vLLM 中评估准确性。

1) 加载模型

使用 AutoModelForCausalLM 加载模型

from transformers import AutoTokenizer, AutoModelForCausalLM

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

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

2) 应用量化

对于 fp8 量化,我们可以通过简单的 PTQ 量化来恢复准确性。

我们建议使用 FP8_DYNAMIC 方案来量化所有 Linear 层,该方案使用:- 权重的静态、每通道量化 - 激活的动态、每令牌量化。

由于简单的 PTQ 不需要用于权重量化的数据,并且激活是动态量化的,因此我们不需要为该量化流程提供任何校准数据。

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.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-Dynamic"
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)

我们已成功创建一个 fp8 模型!

3) 评估准确性

安装 vllmlm-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 令牌的存在敏感。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 上提出问题