@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 175 lines 5.49 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # flake8: noqa
5 """Tests fp8 models against ground truth generation
6 Note: these tests will only pass on L4 GPU.
7 """
8
9 import pytest
10
11 from tests.quantization.utils import is_quant_method_supported
12 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype
13 from vllm.platforms import current_platform
14 from ..utils import check_logprobs_close
15
16
17 @pytest.mark.skipif(
18 not is_quant_method_supported("fp8"),
19 reason="fp8 is not supported on this GPU type.",
20 )
21 @pytest.mark.parametrize(
22 "kv_cache_dtype,base_model,test_model",
23 [
24 # Test FP8 checkpoint w. fp8_e4m3 kv-cache scaling factors.
25 (
26 "fp8_e4m3",
27 "meta-llama/Llama-3.2-1B-Instruct",
28 "nm-testing/Llama-3.2-1B-Instruct-FP8-KV",
29 ),
30 # Test BF16 checkpoint w. fp8_e5m2 kv-cache.
31 (
32 "fp8_e5m2",
33 "meta-llama/Llama-3.2-1B-Instruct",
34 "meta-llama/Llama-3.2-1B-Instruct",
35 ),
36 # Test BF16 checkpoint w. fp8_e4m3 kv-cache scaling factors in json.
37 (
38 "fp8_e4m3",
39 "meta-llama/Llama-3.2-1B-Instruct",
40 "meta-llama/Llama-3.2-1B-Instruct",
41 ),
42 ],
43 )
44 # Due to low-precision numerical divergence, we only test logprob of 4 tokens
45 @pytest.mark.parametrize("max_tokens", [4])
46 @pytest.mark.parametrize("enforce_eager", [True])
47 @pytest.mark.parametrize("backend", ["FLASH_ATTN"])
48 # NOTE: Increasing this in this suite will fail CI because we currently cannot
49 # reset distributed env properly. Use a value > 1 just when you test.
50 @pytest.mark.parametrize("tensor_parallel_size", [1])
51 def test_models(
52 vllm_runner,
53 example_prompts,
54 kv_cache_dtype: str,
55 base_model: str,
56 test_model: str,
57 max_tokens: int,
58 enforce_eager: bool,
59 backend: str,
60 tensor_parallel_size: int,
61 monkeypatch: pytest.MonkeyPatch,
62 ) -> None:
63 """
64 Only checks log probs match to cover the discrepancy in
65 numerical sensitive kernels.
66 """
67
68 if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
69 pytest.skip(f"{kv_cache_dtype} is currently not supported on ROCm/HIP.")
70 if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
71 pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
72
73 if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
74 pytest.skip(
75 f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
76 )
77
78 with monkeypatch.context() as m:
79 m.setenv("TOKENIZERS_PARALLELISM", "true")
80
81 MAX_MODEL_LEN = 1024
82 NUM_LOG_PROBS = 8
83
84 with vllm_runner(
85 base_model,
86 max_model_len=MAX_MODEL_LEN,
87 tensor_parallel_size=tensor_parallel_size,
88 enforce_eager=enforce_eager,
89 kv_cache_dtype="auto",
90 attention_config={"backend": backend},
91 ) as vllm_model:
92 baseline_outputs = vllm_model.generate_greedy_logprobs(
93 example_prompts, max_tokens, NUM_LOG_PROBS
94 )
95
96 with vllm_runner(
97 test_model,
98 max_model_len=MAX_MODEL_LEN,
99 tensor_parallel_size=tensor_parallel_size,
100 enforce_eager=enforce_eager,
101 kv_cache_dtype=kv_cache_dtype,
102 attention_config={"backend": backend},
103 ) as vllm_model:
104 test_outputs = vllm_model.generate_greedy_logprobs(
105 example_prompts, max_tokens, NUM_LOG_PROBS
106 )
107
108 check_logprobs_close(
109 outputs_0_lst=baseline_outputs,
110 outputs_1_lst=test_outputs,
111 name_0="fp16_kv_cache",
112 name_1="fp8_kv_cache",
113 )
114
115
116 @pytest.mark.cpu_model
117 @pytest.mark.skipif(not current_platform.is_cpu(), reason="test for the CPU backend.")
118 @pytest.mark.parametrize(
119 "kv_cache_dtype,base_model,test_model",
120 [
121 # Test BF16 checkpoint w. fp8_e5m2 kv-cache.
122 (
123 "fp8_e5m2",
124 "meta-llama/Llama-3.2-1B-Instruct",
125 "meta-llama/Llama-3.2-1B-Instruct",
126 ),
127 ],
128 )
129 # Due to low-precision numerical divergence, we only test logprob of 4 tokens
130 @pytest.mark.parametrize("max_tokens", [4])
131 def test_cpu_models(
132 vllm_runner,
133 example_prompts,
134 kv_cache_dtype: str,
135 base_model: str,
136 test_model: str,
137 max_tokens: int,
138 monkeypatch: pytest.MonkeyPatch,
139 ) -> None:
140 """
141 Only checks log probs match to cover the discrepancy in
142 numerical sensitive kernels.
143 """
144 with monkeypatch.context() as m:
145 m.setenv("TOKENIZERS_PARALLELISM", "true")
146
147 MAX_MODEL_LEN = 1024
148 NUM_LOG_PROBS = 8
149
150 with vllm_runner(
151 base_model,
152 max_model_len=MAX_MODEL_LEN,
153 dtype="bfloat16",
154 kv_cache_dtype="auto",
155 ) as vllm_model:
156 baseline_outputs = vllm_model.generate_greedy_logprobs(
157 example_prompts, max_tokens, NUM_LOG_PROBS
158 )
159
160 with vllm_runner(
161 test_model,
162 max_model_len=MAX_MODEL_LEN,
163 dtype="bfloat16",
164 kv_cache_dtype=kv_cache_dtype,
165 ) as vllm_model:
166 test_outputs = vllm_model.generate_greedy_logprobs(
167 example_prompts, max_tokens, NUM_LOG_PROBS
168 )
169
170 check_logprobs_close(
171 outputs_0_lst=baseline_outputs,
172 outputs_1_lst=test_outputs,
173 name_0="bf16_kv_cache",
174 name_1="fp8_kv_cache",
175 )