@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 135 lines 5.03 KB
Rendered Raw
1 # NVIDIA Model Optimizer
2
3 The [NVIDIA Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) is a library designed to optimize models for inference with NVIDIA GPUs. It includes tools for Post-Training Quantization (PTQ) and Quantization Aware Training (QAT) of Large Language Models (LLMs), Vision Language Models (VLMs), and diffusion models.
4
5 We recommend installing the library with:
6
7 ```bash
8 pip install nvidia-modelopt
9 ```
10
11 ## Supported ModelOpt checkpoint formats
12
13 vLLM detects ModelOpt checkpoints via `hf_quant_config.json` and supports the
14 following `quantization.quant_algo` values:
15
16 - `FP8`: per-tensor weight scale (+ optional static activation scale).
17 - `FP8_PER_CHANNEL_PER_TOKEN`: per-channel weight scale and dynamic per-token activation quantization.
18 - `FP8_PB_WO` (ModelOpt may emit `fp8_pb_wo`): block-scaled FP8 weight-only (typically 128×128 blocks).
19 - `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`).
20 - `W4A16_NVFP4`: weight-only ModelOpt NVFP4 checkpoints with 16-bit
21 activations (use `quantization="modelopt_fp4"`).
22 - `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`).
23
24 !!! note
25 For NVFP4 checkpoints, vLLM selects a GEMM kernel automatically at load
26 time from the backends available on the current platform (CUTLASS,
27 FlashInfer, Marlin, and others). On GPUs without a supported native FP4
28 GEMM kernel, vLLM falls back to weight-only (W4A16) execution via Marlin
29 and logs a warning; this may reduce throughput for compute-heavy
30 workloads. Use `--linear-backend` to override the automatic selection
31 (this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment
32 variable). Values relevant to NVFP4 include `cutlass`,
33 `flashinfer_cutlass`, `flashinfer_cutedsl`, `flashinfer_trtllm`,
34 `flashinfer_cudnn`, and `marlin`; see `KernelConfig` on the
35 [Engine Arguments](../../configuration/engine_args.md) page and shown by
36 `vllm serve --help=KernelConfig`.
37
38 !!! note
39 For models quantized to MXFP8 with BF16 activations on SM100-family GPUs,
40 use `--linear-backend flashinfer_trtllm` to select FlashInfer's TensorRT-LLM
41 GEMM backend.
42
43 ## Quantizing HuggingFace Models with PTQ
44
45 You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory.
46
47 Below is an example showing how to quantize a model using modelopt's PTQ API:
48
49 ??? code
50
51 ```python
52 import modelopt.torch.quantization as mtq
53 from transformers import AutoModelForCausalLM
54
55 # Load the model from HuggingFace
56 model = AutoModelForCausalLM.from_pretrained("<path_or_model_id>")
57
58 # Select the quantization config, for example, FP8
59 config = mtq.FP8_DEFAULT_CFG
60
61 # Define a forward loop function for calibration
62 def forward_loop(model):
63 for data in calib_set:
64 model(data)
65
66 # PTQ with in-place replacement of quantized modules
67 model = mtq.quantize(model, config, forward_loop)
68 ```
69
70 After the model is quantized, you can export it to a quantized checkpoint using the export API:
71
72 ```python
73 import torch
74 from modelopt.torch.export import export_hf_checkpoint
75
76 with torch.inference_mode():
77 export_hf_checkpoint(
78 model, # The quantized model.
79 export_dir, # The directory where the exported files will be stored.
80 )
81 ```
82
83 The quantized checkpoint can then be deployed with vLLM. As an example, the following code shows how to deploy `nvidia/Llama-3.1-8B-Instruct-FP8`, which is the FP8 quantized checkpoint derived from `meta-llama/Llama-3.1-8B-Instruct`, using vLLM:
84
85 ??? code
86
87 ```python
88 from vllm import LLM, SamplingParams
89
90 def main():
91 model_id = "nvidia/Llama-3.1-8B-Instruct-FP8"
92
93 # Ensure you specify quantization="modelopt" when loading the modelopt checkpoint
94 llm = LLM(model=model_id, quantization="modelopt", trust_remote_code=True)
95
96 sampling_params = SamplingParams(temperature=0.8, top_p=0.9)
97
98 prompts = [
99 "Hello, my name is",
100 "The president of the United States is",
101 "The capital of France is",
102 "The future of AI is",
103 ]
104
105 outputs = llm.generate(prompts, sampling_params)
106
107 for output in outputs:
108 prompt = output.prompt
109 generated_text = output.outputs[0].text
110 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
111
112 if __name__ == "__main__":
113 main()
114 ```
115
116 ## Running the OpenAI-compatible server
117
118 To serve a local ModelOpt checkpoint via the OpenAI-compatible API:
119
120 ```bash
121 vllm serve <path_to_exported_checkpoint> \
122 --quantization modelopt \
123 --host 0.0.0.0 --port 8000
124 ```
125
126 ## Testing (local checkpoints)
127
128 vLLM's ModelOpt unit tests are gated by local checkpoint paths and are skipped
129 by default in CI. To run the tests locally:
130
131 ```bash
132 export VLLM_TEST_MODELOPT_FP8_PC_PT_MODEL_PATH=<path_to_fp8_pc_pt_checkpoint>
133 export VLLM_TEST_MODELOPT_FP8_PB_WO_MODEL_PATH=<path_to_fp8_pb_wo_checkpoint>
134 pytest -q tests/quantization/test_modelopt.py
135 ```