[CI] Move distributed small LM eval to B200 (#47048)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lucas Wilkinson committed
Jun 30, 2026 at 13:34 UTC
27d5f78b638d385904b7409ca398272fc08eead7
4 files changed
+68
-14
.buildkite/test_areas/lm_eval.yaml
+6
-5
@@ -54,8 +54,8 @@ steps:
54
- export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
55
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
56
57
-- label: LM Eval Small Models (2xB200)
58
- key: lm-eval-small-models-2xb200
57
+- label: LM Eval Small Models (1xB200)
58
+ key: lm-eval-small-models-1xb200
59
timeout_in_minutes: 120
60
device: b200-k8s
61
optional: true
@@ -65,9 +65,10 @@ steps:
65
commands:
66
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
67
68
-- label: LM Eval Small Models (2xL4)
69
- key: lm-eval-small-models-tp
70
- timeout_in_minutes: 10
68
+- label: LM Eval Small Models Distributed (2xB200)
69
+ key: lm-eval-small-models-distributed-2xb200
70
+ timeout_in_minutes: 120
71
+ device: b200-k8s
72
num_devices: 2
73
optional: true
74
source_file_dependencies:
tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml
+1
@@ -3,6 +3,7 @@ accuracy_threshold: 0.84
3
num_questions: 1319
4
num_fewshot: 5
5
startup_max_wait_seconds: 1200
6
+use_chat_completions: true
7
server_args: >-
8
--enforce-eager
9
--max-model-len 4096
tests/evals/gsm8k/gsm8k_eval.py
+59
-9
@@ -110,6 +110,39 @@ async def call_vllm_api(
110
return "", 0
111
112
113
+async def call_vllm_chat_api(
114
+ session: aiohttp.ClientSession,
115
+ model: str,
116
+ prompt: str,
117
+ temperature: float,
118
+ max_tokens: int,
119
+ stop: list[str] | None = None,
120
+ url: str | None = None,
121
+ seed: int | None = None,
122
+) -> tuple[str, int]:
123
+ """Call vLLM's OpenAI-compatible chat completions endpoint."""
124
+ data = {
125
+ "model": model,
126
+ "messages": [{"role": "user", "content": prompt}],
127
+ "temperature": temperature,
128
+ "max_tokens": max_tokens,
129
+ "stop": stop,
130
+ }
131
+ if seed is not None:
132
+ data["seed"] = seed
133
+
134
+ try:
135
+ async with session.post(f"{url}/v1/chat/completions", json=data) as response:
136
+ response.raise_for_status()
137
+ result = await response.json()
138
+ text = result["choices"][0]["message"]["content"] or ""
139
+ completion_tokens = result.get("usage", {}).get("completion_tokens", 0)
140
+ return text, completion_tokens
141
+ except Exception as e:
142
+ print(f"Error calling vLLM chat API ({type(e).__name__}): {e}")
143
+ return "", 0
144
+
145
+
146
def _build_gsm8k_prompts(
147
num_questions: int = 1319,
148
num_shots: int = 5,
@@ -173,6 +206,8 @@ def evaluate_gsm8k(
206
num_questions: int = 1319,
207
num_shots: int = 5,
208
max_tokens: int = 256,
209
+ model: str | None = None,
210
+ use_chat_completions: bool = False,
211
host: str = "http://127.0.0.1",
212
port: int = 8000,
213
temperature: float = 0.0,
@@ -193,15 +228,30 @@ def evaluate_gsm8k(
228
output_tokens: list[int] = [0] * num_questions
229
230
async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]:
196
- answer, tokens = await call_vllm_api(
197
- session=session,
198
- prompt=prompts[i],
199
- temperature=temperature,
200
- max_tokens=max_tokens,
201
- stop=["Question", "Assistant:", "<|separator|>"],
202
- url=base_url,
203
- seed=seed,
204
- )
231
+ stop = ["Question", "Assistant:", "<|separator|>"]
232
+ if use_chat_completions:
233
+ if model is None:
234
+ raise ValueError("model is required for chat completions")
235
+ answer, tokens = await call_vllm_chat_api(
236
+ session=session,
237
+ model=model,
238
+ prompt=prompts[i],
239
+ temperature=temperature,
240
+ max_tokens=max_tokens,
241
+ stop=stop,
242
+ url=base_url,
243
+ seed=seed,
244
+ )
245
+ else:
246
+ answer, tokens = await call_vllm_api(
247
+ session=session,
248
+ prompt=prompts[i],
249
+ temperature=temperature,
250
+ max_tokens=max_tokens,
251
+ stop=stop,
252
+ url=base_url,
253
+ seed=seed,
254
+ )
255
states[i] = answer
256
output_tokens[i] = tokens
257
return answer, tokens
tests/evals/gsm8k/test_gsm8k_correctness.py
+2
@@ -65,6 +65,8 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
65
num_questions=eval_config["num_questions"],
66
num_shots=eval_config["num_fewshot"],
67
max_tokens=eval_config.get("max_tokens", 256),
68
+ model=eval_config["model_name"],
69
+ use_chat_completions=eval_config.get("use_chat_completions", False),
70
host=host,
71
port=port,
72
request_timeout_seconds=request_timeout_seconds,