NVIDIA TensorRT 模型最佳化器

NVIDIA TensorRT 模型最佳化器 是一個旨在最佳化模型以在 NVIDIA GPU 上進行推理的庫。它包括用於大語言模型 (LLM)、視覺語言模型 (VLM) 和擴散模型的訓練後量化 (PTQ) 和量化感知訓練 (QAT) 的工具。

我們建議透過以下方式安裝該庫：

pip install nvidia-modelopt

使用 PTQ 量化 HuggingFace 模型

您可以使用 TensorRT 模型最佳化器倉庫中提供的示例指令碼來量化 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()