@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 167 lines 5.62 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """HTTP-based batch invariance test: send requests to a running
4 vLLM server and compare BS=1 vs BS=N results (tokens and per-step logprobs).
5
6 Environment variables:
7 - VLLM_TEST_MODEL: served model name (e.g., Qwen/Qwen3-1.7B / DeepSeek-R1)
8 - VLLM_TP_SIZE: tensor parallelism size (e.g., 4)
9
10 """
11
12 import os
13 import random
14 import sys
15 from typing import Any
16
17 import openai
18 import pytest
19 from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda
20
21 from tests.utils import RemoteOpenAIServer
22
23
24 def _request_completion(
25 client: openai.OpenAI,
26 model: str,
27 prompt: Any,
28 sp: dict[str, Any],
29 max_retries: int = 3,
30 retry_backoff: float = 0.5,
31 ) -> dict[str, Any] | None:
32 payload: dict[str, Any] = {"model": model, "prompt": prompt}
33 payload.update(sp)
34
35 for attempt in range(max_retries + 1):
36 try:
37 completion = client.completions.create(**payload)
38 # Convert to plain dict so downstream logic can keep using
39 # dict-style access just like with raw HTTP JSON.
40 return completion.model_dump()
41 except Exception as e: # pragma: no cover
42 if attempt < max_retries:
43 import time as _t
44
45 _t.sleep(retry_backoff * (2**attempt))
46 continue
47 sys.stderr.write(f"Error: {e}\n")
48 return None
49 return None
50
51
52 def _extract_tokens_and_logprobs(
53 choice: dict[str, Any],
54 ) -> tuple[list[Any], list[float] | None]:
55 tokens: list[Any] = []
56 token_logprobs: list[float] | None = None
57 lp = choice.get("logprobs")
58 if lp and isinstance(lp, dict):
59 tokens = lp.get("token_ids") or lp.get("tokens") or []
60 token_logprobs = lp.get("token_logprobs", None)
61 return tokens, token_logprobs
62
63
64 def _compare_bs1_vs_bsn_single_process(
65 prompts: list[str],
66 sp_kwargs: dict[str, Any],
67 client: openai.OpenAI,
68 model_name: str,
69 ) -> None:
70 # BS=1
71 bs1_tokens_per_prompt: list[list[Any]] = []
72 bs1_logprobs_per_prompt: list[list[float] | None] = []
73 for p in prompts:
74 resp = _request_completion(client, model_name, p, sp_kwargs)
75 if resp is None or not resp.get("choices"):
76 raise AssertionError("BS=1 empty/failed response")
77 choice = resp["choices"][0]
78 toks, lps = _extract_tokens_and_logprobs(choice)
79 if lps is None:
80 raise AssertionError(
81 "logprobs not returned; ensure server supports 'logprobs'"
82 )
83 bs1_tokens_per_prompt.append(list(toks))
84 bs1_logprobs_per_prompt.append(list(lps))
85
86 # BS=N
87 bsN_tokens_per_prompt: list[list[Any]] = [None] * len(prompts) # type: ignore[list-item]
88 bsN_logprobs_per_prompt: list[list[float] | None] = [None] * len(prompts)
89 resp = _request_completion(client, model_name, prompts, sp_kwargs)
90 if resp is None or not resp.get("choices"):
91 raise AssertionError("BS=N empty/failed batched response")
92 choices = resp.get("choices", [])
93 if len(choices) != len(prompts):
94 raise AssertionError(
95 f"BS=N choices length {len(choices)} != num prompts {len(prompts)}"
96 )
97 for idx, choice in enumerate(choices):
98 toks, lps = _extract_tokens_and_logprobs(choice)
99 if lps is None:
100 raise AssertionError(f"BS=N missing logprobs for prompt {idx}")
101 bsN_tokens_per_prompt[idx] = list(toks)
102 bsN_logprobs_per_prompt[idx] = list(lps)
103
104 # compare
105 for i, (tokens_bs1, tokens_bsN, logprobs_bs1, logprobs_bsN) in enumerate(
106 zip(
107 bs1_tokens_per_prompt,
108 bsN_tokens_per_prompt,
109 bs1_logprobs_per_prompt,
110 bsN_logprobs_per_prompt,
111 )
112 ):
113 if tokens_bs1 != tokens_bsN:
114 raise AssertionError(
115 f"Prompt {i} (sampling): Different tokens sampled. "
116 f"BS=1 tokens: {tokens_bs1} BS=N tokens: {tokens_bsN}"
117 )
118 if logprobs_bs1 is None or logprobs_bsN is None:
119 raise AssertionError(f"Prompt {i}: Missing logprobs in one of the runs")
120 if len(logprobs_bs1) != len(logprobs_bsN):
121 raise AssertionError(
122 f"Prompt {i}: Different number of steps: "
123 f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bsN)} (BS=N)."
124 )
125 for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):
126 if a != b:
127 diff = abs(a - b)
128 raise AssertionError(
129 f"Prompt {i} Step {t}: Bitwise mismatch "
130 f"(abs diff={diff:.6e}). "
131 f"BS=1 tokens: {tokens_bs1} BS=N tokens: {tokens_bsN}"
132 )
133
134
135 @skip_if_not_cuda
136 @pytest.mark.parametrize("backend", BACKENDS)
137 def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
138 backend: str,
139 ) -> None:
140 random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
141 prompts_all = [_random_prompt(10, 50) for _ in range(32)]
142
143 sp_kwargs: dict[str, Any] = {
144 "temperature": 0.6,
145 "top_p": 1.0,
146 "max_tokens": 8,
147 "seed": 42,
148 "logprobs": 5,
149 }
150
151 tp_size = os.getenv("VLLM_TP_SIZE", "1")
152 server_args: list[str] = [
153 "--max-model-len=8192",
154 "--max-num-seqs=32",
155 f"--attention-backend={backend}",
156 ]
157 if tp_size:
158 server_args += ["-tp", tp_size]
159
160 with RemoteOpenAIServer(TEST_MODEL, server_args) as server:
161 client = server.get_client()
162 _compare_bs1_vs_bsn_single_process(
163 prompts=prompts_all,
164 sp_kwargs=sp_kwargs,
165 client=client,
166 model_name=TEST_MODEL,
167 )