main
py 153 lines 5.14 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # flake8: noqa
4 """Tests Model Optimizer nvfp4 models against ground truth generation
5 Note: these tests will only pass on B200
6 """
7
8 import os
9 from typing import List
10
11 import pytest
12 from transformers import AutoTokenizer
13
14 from tests.quantization.utils import is_quant_method_supported
15 from vllm import LLM, SamplingParams
16
17 from vllm.platforms import current_platform
18
19 os.environ["TOKENIZERS_PARALLELISM"] = "true"
20
21 MAX_MODEL_LEN = 1024
22
23 MODELS = ["nvidia/Llama-3.3-70B-Instruct-FP4"]
24
25 EXPECTED_STRS_MAP = {
26 "nvidia/Llama-3.3-70B-Instruct-FP4": [
27 "vLLM (Vectorized Large Language Model) is indeed a high-throughput and memory-efficient inference",
28 "Here are the major milestones in the development of artificial intelligence (AI) from 1950 to ",
29 "Artificial intelligence (AI) and human intelligence (HI) are two distinct forms of intelligence that process",
30 "A neural network is a type of machine learning model inspired by the structure and function of the human brain",
31 "In the heart of a cutting-edge robotics lab, a team of engineers had been working tirelessly to push",
32 "The COVID-19 pandemic has had a profound impact on global economic structures and future business models, leading",
33 "The Mona Lisa, painted by Leonardo da Vinci in the early 16th century, is one of",
34 "Here are the translations:\n\n* Japanese: (Sasuga no tori ga miwa o ts",
35 ]
36 }
37
38
39 # This test compares against golden strings for exact match since
40 # there is no baseline implementation to compare against
41 # and is unstable w.r.t specifics of the fp4 implementation or
42 # the hardware being run on.
43 # Disabled to prevent it from breaking the build
44 @pytest.mark.skip(
45 reason="Prevent unstable test based on golden strings from breaking the build "
46 " and test input model being too large and hanging the system."
47 )
48 @pytest.mark.skipif(
49 not is_quant_method_supported("modelopt_fp4"),
50 reason="modelopt_fp4 is not supported on this GPU type.",
51 )
52 @pytest.mark.parametrize("model_name", MODELS)
53 def test_models(example_prompts, model_name) -> None:
54 llm = LLM(
55 model=model_name,
56 max_model_len=MAX_MODEL_LEN,
57 trust_remote_code=True,
58 enforce_eager=True,
59 quantization="modelopt_fp4",
60 )
61
62 tokenizer = AutoTokenizer.from_pretrained(model_name)
63 formatted_prompts = [
64 tokenizer.apply_chat_template(
65 [{"role": "user", "content": prompt}],
66 tokenize=False,
67 add_generation_prompt=True,
68 )
69 for prompt in example_prompts
70 ]
71 params = SamplingParams(max_tokens=20, temperature=0)
72 generations: List[str] = []
73 # Note: these need to be run 1 at a time due to numerical precision,
74 # since the expected strs were generated this way.
75 for prompt in formatted_prompts:
76 outputs = llm.generate(prompt, params)
77 generations.append(outputs[0].outputs[0].text)
78 del llm
79
80 print(model_name, generations)
81 expected_strs = EXPECTED_STRS_MAP[model_name]
82 for i in range(len(example_prompts)):
83 generated_str = generations[i]
84 expected_str = expected_strs[i]
85 assert expected_str == generated_str, (
86 f"Test{i}:\nExpected: {expected_str!r}\nvLLM: {generated_str!r}"
87 )
88
89
90 EAGER = [True, False]
91
92 SM_100_NVFP4_BACKENDS = [
93 "flashinfer_cutedsl",
94 "flashinfer_cudnn",
95 "flashinfer_trtllm",
96 "flashinfer_cutlass",
97 "torch",
98 ]
99
100
101 @pytest.mark.parametrize("model", ["nvidia/Llama-3.1-8B-Instruct-NVFP4"])
102 @pytest.mark.parametrize("eager", EAGER)
103 @pytest.mark.parametrize(
104 "backend",
105 [
106 "emulation",
107 "flashinfer_cutedsl",
108 "flashinfer_cudnn",
109 "flashinfer_trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used
110 "flashinfer_cutlass",
111 "torch",
112 ],
113 )
114 def test_nvfp4(vllm_runner, model, eager, backend):
115 if backend == "flashinfer_cutedsl" and not (
116 current_platform.is_device_capability_family(100)
117 ):
118 pytest.skip("The flashinfer_cutedsl backend is only supported on SM10x")
119
120 if (
121 not current_platform.has_device_capability(100)
122 and backend in SM_100_NVFP4_BACKENDS
123 ):
124 pytest.skip(
125 f"The backend {backend} is not supported with current_platform.has_device_capability(100) == False"
126 )
127
128 with vllm_runner(model, enforce_eager=eager, linear_backend=backend) as llm:
129 output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)
130 assert output[0][1] == "1 2 3 4 5 6"
131
132
133 @pytest.mark.parametrize(
134 "model",
135 [
136 "nvidia/Qwen3-30B-A3B-NVFP4",
137 "RedHatAI/Qwen3-30B-A3B-NVFP4",
138 ],
139 )
140 @pytest.mark.parametrize("backend", ["emulation"])
141 @pytest.mark.skipif(
142 not current_platform.is_rocm(),
143 reason="NVFP4 MOE emulation is only useful on AMD Instinct MI3xx",
144 )
145 def test_nvfp4_moe(vllm_runner, model, backend):
146 with vllm_runner(
147 model,
148 moe_backend=backend,
149 linear_backend=backend,
150 load_format="dummy",
151 hf_overrides={"num_hidden_layers": 2},
152 ) as llm:
153 _ = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)