跳到内容

fp4 量化

llm-compressor 支持将权重和激活量化到 fp4,以节省内存并使用 vLLM 加速推理。特别是,支持 nvfp4 - NVIDIA Blackwell GPU 架构引入的一种 4 位浮点编码格式。

安装

首先,安装

git clone https://github.com/vllm-project/llm-compressor.git
cd llm-compressor
pip install -e .

快速入门

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

python3 llama3_example.py

生成的模型 Meta-Llama-3-8B-Instruct-NVFP4 已准备好加载到 vLLM 中。注意:如果在 SM100 以下的机器上进行推理,vLLM 将不会激活量化,仅进行权重量化。

代码演练

现在,我们将逐步介绍示例中的代码:1) 加载模型 2) 准备校准数据 3) 应用量化

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) 准备校准数据

准备校准数据。nvfp4 量化会为权重生成每张量全局尺度和每组(大小为 16)局部量化尺度,以及为激活生成每张量全局尺度。局部激活量化尺度将在推理时动态生成。我们需要一些样本数据来校准全局激活尺度。通常,少量样本就足够了。在此示例中,我们使用 20 个样本。

使用与部署时使用的数据类型非常匹配的校准数据很有用。如果您对模型进行了微调,使用训练数据样本是个好主意。在我们的案例中,我们正在量化一个指令调整过的通用模型,因此我们将使用 ultrachat 数据集。

3) 应用量化

数据集准备好后,我们现在将应用量化。

我们首先选择量化算法。

在我们的案例中,我们将把 nvfp4 的默认 QuantizationModifier 配方应用于所有线性层。

有关创建复杂配方的更多信息,请参阅 Recipes 文档。

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

# Configure the quantization algorithm to run.
recipe = QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=["lm_head"])

# Apply quantization.
oneshot(
    model=model,
    dataset=ds,
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)

# Save to disk compressed.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

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

量化 MoEs

为了量化 MoEs,MoE 校准现在由管道自动处理。可以在 llama4_example.py 中找到量化 Llama4 的示例。管道会自动应用适当的 MoE 校准上下文,该上下文

  1. 使模型线性化,以便在 vLLM 中进行量化和执行。这是必需的,因为原生模型定义在其 MoE 块中不包含 torch.nn.Linear 层,而这是 LLM Compressor 运行量化的要求。
  2. 确保专家被正确量化,因为并非所有专家在校准期间都会被激活。

同样,可以在 qwen_30b_a3b.py 中找到量化 Qwen3-30B-A3B 模型的示例。该模型使用上下文 MoE 校准,该校准会临时更新模型定义以使用 Qwen3MoeSparseMoeBlock,从而在校准期间更新 MoE 块中前向传递的处理方式。您可以随意在 llm-compressor/src/llmcompressor/modeling/qwen3_moe.py 下更新定义,以尝试这种行为并评估其对量化性能的影响。