@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 130 lines 3.43 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """E2E tests for GGUF plugin functionality."""
4
5 import os
6 from typing import NamedTuple
7
8 import pytest
9 from transformers import AutoTokenizer
10
11 from ...conftest import VllmRunner
12 from ...models.utils import check_logprobs_close
13 from ...utils import multi_gpu_test
14
15 os.environ["TOKENIZERS_PARALLELISM"] = "true"
16
17 MAX_MODEL_LEN = 1024
18
19
20 class GGUFTestConfig(NamedTuple):
21 original_model: str
22 gguf_model_path: str # Full path to .gguf file
23
24
25 QWEN3_CONFIG = GGUFTestConfig(
26 original_model="Qwen/Qwen3-0.6B",
27 gguf_model_path="unsloth/Qwen3-0.6B-GGUF:Q8_0",
28 )
29
30
31 OLMOE_CONFIG = GGUFTestConfig(
32 original_model="allenai/OLMoE-1B-7B-0125",
33 gguf_model_path="allenai/OLMoE-1B-7B-0125-GGUF:Q6_K",
34 )
35
36
37 MODELS = [
38 QWEN3_CONFIG,
39 OLMOE_CONFIG,
40 ]
41
42
43 def check_model_outputs(
44 vllm_runner: type[VllmRunner],
45 prompts: list[str],
46 model: GGUFTestConfig,
47 dtype: str,
48 max_tokens: int,
49 num_logprobs: int,
50 tp_size: int,
51 ):
52 tokenizer = AutoTokenizer.from_pretrained(model.original_model)
53 if tokenizer.chat_template is not None:
54 messages = [[{"role": "user", "content": prompt}] for prompt in prompts]
55 prompts = tokenizer.apply_chat_template(
56 messages, tokenize=False, add_generation_prompt=True
57 )
58
59 # Run gguf model.
60 with vllm_runner(
61 model_name=model.gguf_model_path,
62 enforce_eager=True,
63 tokenizer_name=model.original_model,
64 dtype=dtype,
65 max_model_len=MAX_MODEL_LEN,
66 tensor_parallel_size=tp_size,
67 ) as gguf_model:
68 gguf_outputs = gguf_model.generate_greedy_logprobs(
69 prompts[:-1], max_tokens, num_logprobs
70 )
71
72 # Run unquantized model.
73 # Should run with tp=1, otherwise the test will stuck at
74 # nccl initialization.
75 with vllm_runner(
76 model_name=model.original_model,
77 enforce_eager=True, # faster tests
78 dtype=dtype,
79 max_model_len=MAX_MODEL_LEN,
80 tensor_parallel_size=1,
81 ) as original_model:
82 original_outputs = original_model.generate_greedy_logprobs(
83 prompts[:-1], max_tokens, num_logprobs
84 )
85
86 check_logprobs_close(
87 outputs_0_lst=original_outputs,
88 outputs_1_lst=gguf_outputs,
89 name_0="original",
90 name_1="gguf",
91 )
92
93
94 @pytest.mark.parametrize("model", MODELS)
95 @pytest.mark.parametrize("dtype", ["bfloat16"])
96 @pytest.mark.parametrize("max_tokens", [32])
97 @pytest.mark.parametrize("num_logprobs", [8])
98 @pytest.mark.parametrize("tp_size", [1])
99 def test_models(
100 vllm_runner: type[VllmRunner],
101 example_prompts: list[str],
102 model: GGUFTestConfig,
103 dtype: str,
104 max_tokens: int,
105 num_logprobs: int,
106 tp_size: int,
107 ) -> None:
108 check_model_outputs(
109 vllm_runner, example_prompts, model, dtype, max_tokens, num_logprobs, tp_size
110 )
111
112
113 @pytest.mark.parametrize("model", MODELS)
114 @pytest.mark.parametrize("dtype", ["half"])
115 @pytest.mark.parametrize("max_tokens", [8])
116 @pytest.mark.parametrize("num_logprobs", [5])
117 @pytest.mark.parametrize("tp_size", [2])
118 @multi_gpu_test(num_gpus=2)
119 def test_distributed(
120 vllm_runner: type[VllmRunner],
121 example_prompts: list[str],
122 model: GGUFTestConfig,
123 dtype: str,
124 max_tokens: int,
125 num_logprobs: int,
126 tp_size: int,
127 ) -> None:
128 check_model_outputs(
129 vllm_runner, example_prompts, model, dtype, max_tokens, num_logprobs, tp_size
130 )