[DBO][CI] Increase the coverage of prefill DBO in test_dbo.py (#48628)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Aug 18, 2026 at 12:00 UTC
9842d701450214d4b78cd9aefb8eee0c616bce33
1 file changed
+20
-11
tests/v1/distributed/test_dbo.py
+20
-11
@@ -11,7 +11,6 @@ correctly with the DeepSeek-V2-Lite model using GSM8K evaluation.
11
import pytest
12
import torch
13
14
-from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k
14
from tests.utils import RemoteOpenAIServer
15
from vllm.utils.import_utils import has_deep_ep
16
@@ -58,6 +57,7 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
57
"""
58
Test DBO with DP+EP using GSM8K evaluation.
59
"""
60
+ lm_eval = pytest.importorskip("lm_eval")
61
required_gpus = DP_SIZE
62
63
if num_gpus_available < required_gpus:
@@ -79,7 +79,7 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
79
"--dbo-decode-token-threshold",
80
"16",
81
"--dbo-prefill-token-threshold",
82
- "256",
82
+ "32",
83
"--all2all-backend",
84
all2all_backend,
85
]
@@ -90,19 +90,28 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
90
max_wait_seconds=600, # Allow time for model loading with DP+EP
91
) as remote_server:
92
# Use host and port directly from RemoteOpenAIServer
93
- host = f"http://{remote_server.host}"
94
- port = remote_server.port
93
+
94
+ base_url = f"http://{remote_server.host}:{remote_server.port}/v1/completions"
95
96
# Run GSM8K evaluation
97
- results = evaluate_gsm8k(
98
- num_questions=NUM_QUESTIONS,
99
- num_shots=NUM_SHOTS,
100
- host=host,
101
- port=port,
97
+ results = lm_eval.simple_evaluate(
98
+ model="local-completions",
99
+ model_args=(
100
+ f"pretrained={MODEL_NAME},"
101
+ f"base_url={base_url},"
102
+ "num_concurrent=512,max_retries=3"
103
+ ),
104
+ tasks=["gsm8k"],
105
+ num_fewshot=NUM_SHOTS,
106
+ limit=NUM_QUESTIONS,
107
)
103
-
108
# Validate accuracy is reasonable
105
- accuracy = results["accuracy"]
109
+ gsm8k = results["results"]["gsm8k"]
110
+ accuracy = gsm8k.get(
111
+ "exact_match,strict-match",
112
+ gsm8k.get("exact_match,flexible-extract"),
113
+ )
114
+ assert accuracy is not None, f"gsm8k exact_match missing: {gsm8k}"
115
assert accuracy >= MIN_ACCURACY, (
116
f"DBO+DP+EP accuracy too low ({all2all_backend}): "
117
f"{accuracy:.3f} < {MIN_ACCURACY:.3f} "