vllm源码剖析11-vLLM 量化模型实践

文章目录

• • [一 vLLM 量化功能总结](#一 vLLM 量化功能总结)
  • • [1.1 量化格式和量化粒度](#1.1 量化格式和量化粒度)
    • [1.2 FP8 格式说明](#1.2 FP8 格式说明)
  • [二 vLLM 量化推理实践](#二 vLLM 量化推理实践)
  • • [2.1 w4a16 的量化推理实践](#2.1 w4a16 的量化推理实践)
    • [2.2 w8a8 权重量化转换实践](#2.2 w8a8 权重量化转换实践)
  • [三 vLLM 的量化模块源码解析](#三 vLLM 的量化模块源码解析)
  • • [3.1 quantization 模块概述](#3.1 quantization 模块概述)
    • [3.2 量化的端到端流程](#3.2 量化的端到端流程)
  • [四 参考资料](#四 参考资料)

一 vLLM 量化功能总结

1.1 量化格式和量化粒度

1. 在 vLLM 的定点整型及 FP8 量化中，支持以下几种常见格式：
   a. W4A16：权重为 4-bit，激活保持 16-bit，例如 FP16 / BF16。常见于 weight-only 量化路径。
   b. W8A16：权重为 8-bit，激活保持 16-bit。这里的 8-bit 权重不只可能是 INT8，在 compressed-tensors 的 FP8 W8A16 路径中也可以是 FP8 权重。
   c. W8A8：权重和激活都使用 8-bit 量化，可以是 INT8，也可以是 FP8，具体取决于量化 scheme。
   d. FP8：vLLM 线性层 / 矩阵乘路径中常见的是 E4M3 相关格式，例如 torch.float8_e4m3fn；在 ROCm 平台上还会处理 torch.float8_e4m3fnuz。

相关 dtype 示例：

f32_type = torch.float32

bf16_type = torch.bfloat16

e4m3_type = torch.float8_e4m3fn

2. 权重、激活的量化粒度总结：

• 权重量化粒度：
  a. Per-channel 量化：通常按输出通道保存 scale；
  b. Per-group 量化：按 group size 分组保存 scale，常见于 INT4 / INT8 weight-only 或 FP4 权重量化；
  c. Per-tensor 量化：整个权重张量或逻辑分片共享量化系数；
  d. Per-block 量化：FP8 W8A8、FP8 W8A16 等路径中可按块大小保存 scale，例如 block-wise FP8 权重量化。
• 激活量化粒度：
  a. Dynamic per-token 量化：每个 token 动态计算 scale，是 W8A8 / FP8 动态激活量化中的常见形式；
  b. Dynamic per-tensor 量化：整个激活张量共享一个 scale，实现更简单；
  c. Per-token-group 量化：FP8 group 量化路径中会按 token 和 group 共同确定 scale；
  d. Static per-tensor 量化：从 checkpoint 中加载固定 input scale，常见于静态激活量化模型。
  另外，要将 INT8/INT4 等量化与 vLLM 一起使用，你需要安装 llm-compressor 库：

pip install llmcompressor==0.8.1

当然，vLLM 支持多种量化方法，例如 awq、gptq、compressed-tensors、bitsandbytes、torchao、inc、online 等。下表展示了 vLLM 中各种量化实现与不同硬件平台的兼容性.

• Volta 对应 SM 7.0 架构，Turing 对应 SM 7.5 架构，Ampere 对应 SM 8.0/8.6 架构，Ada 对应 SM 8.9 架构，Hopper 对应 SM 9.0 架构。
• ✅︎ 表示该量化方法在指定硬件上受支持。
• ❌ 表示该量化方法在指定硬件上不受支持

1.2 FP8 格式说明

深度学习框架中常见的 FP8 表示可以分为 E4M3 和 E5M2。它们的区别主要在指数位和尾数位的分配，因此适合的场景不同。

1. E4M3 使用 1 位符号位、4 位指数位、3 位尾数位。因为尾数位相对更多，它比 E5M2 提供更高的有效精度；但指数位较少，所以动态范围更小。它适合权重存储、激活量化和前向推理中的矩阵乘路径。以 PyTorch 的 torch.float8_e4m3fn 为例，本地环境中：

torch.finfo(torch.float8_e4m3fn).max   # 448.0
torch.finfo(torch.float8_e4m3fn).tiny  # 0.015625

torch.float8_e4m3fn 的命名可以拆开理解：

• float8：总共 8 个比特位；
• e4：Exponent，指数位占 4 位；
• m3：Mantissa，尾数位占 3 位；
• fn：Finite，表示该格式没有 Infinity 编码，保留有限值和 NaN 相关编码。实际量化时，越界值如何处理还取决于具体转换、scale 和 kernel 的 clamp 逻辑。
  在 vLLM 中，FP8 线性方法只支持 float8_e4m3fn，在Nvidia平台默认的 fp8_dtype() 也是 torch.float8_e4m3fn。

class Platform:
    def fp8_dtype(cls) -> torch.dtype:
        """
        Returns the preferred FP8 type on the current platform.

        See the documentation for is_fp8_fnuz for details.
        """
        return torch.float8_e4m3fn

2. E5M2 使用 1 位符号位、5 位指数位、2 位尾数位。它的指数位更多，因此动态范围更大；但尾数位更少，因此精度低于 E4M3。它更常出现在梯度、反向传播或需要更大动态范围的中间值训练相关的场景中。以 PyTorch 的 torch.float8_e5m2 为例，本地环境中：

   torch.finfo(torch.float8_e5m2).max # 57344.0
   torch.finfo(torch.float8_e5m2).tiny # 6.103515625e-05

vLLM 推理量化的主线仍然是 E4M3 相关格式。

二 vLLM 量化推理实践

2.1 w4a16 的量化推理实践

W4A16 指的是权重使用 4-bit 量化，激活仍使用 16-bit 路径计算。在 vLLM 中，AWQ 是常见的 W4A16 量化格式之一，它主要降低权重显存占用。要获取 4 位量化模型，在 Huggingface 或魔塔社区上直接搜索下载 INT4 量化模型。

也可利用 AutoAWQ 或 llmcompressor 工具库，将模型精度从 BF16/FP16 量化至 INT4，从而有效减少整体模型内存占用。这里要区分两个阶段：

1. 量化生成阶段：使用 AutoAWQ 这类工具生成 AWQ checkpoint；
2. vLLM 推理阶段：vLLM 读取量化模型目录中的 quantization_config、quant_config.json 或 quantize_config.json，并选择 AWQ 对应的推理路径。

pip install autoawq
# pip install llmcompressor

安装完 AutoAWQ 后，我们就可以对模型进行量化了。更多详细信息请参阅 AutoAWQ 官方文档。以下是如何量化 mistralai/Mistral-7B-Instruct-v0.2 的示例，当程序结束输出之后就会有一个量化模型的输出文件夹，我在 course10之下运行就会有一个文件夹为code/course10/mistral-instruct-v0.2-awq

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = 'mistralai/Mistral-7B-Instruct-v0.2'
quant_path = 'mistral-instruct-v0.2-awq'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load model
model = AutoAWQForCausalLM.from_pretrained(
    model_path, **{"low_cpu_mem_usage": True, "use_cache": False}
)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

print(f'Model is quantized and saved at "{quant_path}"')

生成 AWQ 模型后，可以用 vLLM 加载它进行推理，我们可以执行以下命令，注意执行时候的模型路径问题。

python code/course10/llm_engine_example.py \
--model code/course10/mistral-instruct-v0.2-awq \
--quantization awq --dtype float16

其中获得的一条结果如下，在命令行输出结果的 RequestOutput.outputs.CompletionOutput.text 记录了完整的输出

2.2 w8a8 权重量化转换实践

三 vLLM 的量化模块源码解析

3.1 quantization 模块概述

3.2 量化的端到端流程

四 参考资料

https://docs.vllm.ai/en/latest/features/quantization/fp8/