跳到内容

NVIDIA Model Optimizer

NVIDIA Model Optimizer 是一个旨在为 NVIDIA GPU 推理优化模型的库。它包含用于大语言模型 (LLM)、视觉语言模型 (VLM) 和扩散模型的训练后量化 (PTQ) 和量化感知训练 (QAT) 工具。

我们建议通过以下命令安装该库

pip install nvidia-modelopt

支持的 ModelOpt 检查点格式

vLLM 通过 hf_quant_config.json 检测 ModelOpt 检查点,并支持以下 quantization.quant_algo

  • FP8:张量级权重缩放(+ 可选的静态激活缩放)。
  • FP8_PER_CHANNEL_PER_TOKEN:通道级权重缩放和动态 Token 级激活量化。
  • FP8_PB_WO (ModelOpt 可能输出 fp8_pb_wo):块缩放的 FP8 仅权重(通常为 128×128 块)。
  • NVFP4:ModelOpt NVFP4 检查点(使用 quantization="modelopt_fp4")。
  • MXFP8:ModelOpt MXFP8 检查点(使用 quantization="modelopt_mxfp8")。

使用 PTQ 量化 HuggingFace 模型

您可以使用 Model Optimizer 存储库中提供的示例脚本来量化 HuggingFace 模型。用于 LLM PTQ 的主要脚本通常位于 examples/llm_ptq 目录下。

以下是一个展示如何使用 modelopt 的 PTQ API 量化模型的示例

代码
import modelopt.torch.quantization as mtq
from transformers import AutoModelForCausalLM

# Load the model from HuggingFace
model = AutoModelForCausalLM.from_pretrained("<path_or_model_id>")

# Select the quantization config, for example, FP8
config = mtq.FP8_DEFAULT_CFG

# Define a forward loop function for calibration
def forward_loop(model):
    for data in calib_set:
        model(data)

# PTQ with in-place replacement of quantized modules
model = mtq.quantize(model, config, forward_loop)

模型量化完成后,您可以使用导出 API 将其导出为量化检查点

import torch
from modelopt.torch.export import export_hf_checkpoint

with torch.inference_mode():
    export_hf_checkpoint(
        model,  # The quantized model.
        export_dir,  # The directory where the exported files will be stored.
    )

然后可以使用 vLLM 部署该量化检查点。作为示例,以下代码展示了如何使用 vLLM 部署 nvidia/Llama-3.1-8B-Instruct-FP8(这是从 meta-llama/Llama-3.1-8B-Instruct 导出的 FP8 量化检查点)

代码
from vllm import LLM, SamplingParams

def main():
    model_id = "nvidia/Llama-3.1-8B-Instruct-FP8"

    # Ensure you specify quantization="modelopt" when loading the modelopt checkpoint
    llm = LLM(model=model_id, quantization="modelopt", trust_remote_code=True)

    sampling_params = SamplingParams(temperature=0.8, top_p=0.9)

    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]

    outputs = llm.generate(prompts, sampling_params)

    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

if __name__ == "__main__":
    main()

运行兼容 OpenAI 的服务器

通过兼容 OpenAI 的 API 为本地 ModelOpt 检查点提供服务

vllm serve <path_to_exported_checkpoint> \
  --quantization modelopt \
  --host 0.0.0.0 --port 8000

测试(本地检查点)

vLLM 的 ModelOpt 单元测试受本地检查点路径限制,在 CI 中默认跳过。要在本地运行测试

export VLLM_TEST_MODELOPT_FP8_PC_PT_MODEL_PATH=<path_to_fp8_pc_pt_checkpoint>
export VLLM_TEST_MODELOPT_FP8_PB_WO_MODEL_PATH=<path_to_fp8_pb_wo_checkpoint>
pytest -q tests/quantization/test_modelopt.py