NVIDIA Model Optimizer¶
NVIDIA Model Optimizer 是一个旨在为 NVIDIA GPU 推理优化模型的库。它包含用于大语言模型 (LLM)、视觉语言模型 (VLM) 和扩散模型的训练后量化 (PTQ) 和量化感知训练 (QAT) 工具。
我们建议通过以下命令安装该库
支持的 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 的 ModelOpt 单元测试受本地检查点路径限制,在 CI 中默认跳过。要在本地运行测试