@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 137 lines 3.79 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from __future__ import annotations
4
5 import lm_eval
6 import pytest
7
8 from tests.utils import large_gpu_mark
9 from vllm.platforms import current_platform
10
11
12 def get_model_args(
13 model_name: str,
14 spec_model_name: str | None,
15 spec_method: str,
16 tp_size: int,
17 model_max_len: int,
18 use_async: bool = True,
19 ) -> dict:
20 speculative_config = {
21 "method": spec_method,
22 "model": spec_model_name,
23 "num_speculative_tokens": 1,
24 "max_model_len": model_max_len,
25 }
26 eplb_config = {
27 "num_redundant_experts": tp_size,
28 "window_size": 128,
29 "step_interval": 1024,
30 "log_balancedness": False,
31 "use_async": use_async,
32 }
33 model_args = {
34 "pretrained": model_name,
35 "dtype": "auto",
36 "add_bos_token": True,
37 "tensor_parallel_size": tp_size,
38 "gpu_memory_utilization": 0.7,
39 "speculative_config": speculative_config,
40 "enable_expert_parallel": True,
41 "eplb_config": eplb_config,
42 "enable_eplb": True,
43 "max_model_len": model_max_len,
44 }
45 return model_args
46
47
48 pytestmark = pytest.mark.skipif(
49 current_platform.is_rocm(),
50 reason="EPLB with Spec Decode is a work in progress on ROCm.",
51 )
52
53
54 @pytest.mark.parametrize(
55 "model_setup",
56 [
57 pytest.param(
58 ("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86),
59 marks=large_gpu_mark(min_gb=80),
60 ),
61 pytest.param(
62 (
63 "eagle",
64 "meta-llama/Llama-4-Scout-17B-16E-Instruct",
65 "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
66 4,
67 0.92,
68 ),
69 marks=pytest.mark.skip(reason="Skipping due to CI OOM issues"),
70 ),
71 ],
72 ids=["qwen3_next_mtp", "llama4_eagle"],
73 )
74 def test_eplb_spec_decode(
75 monkeypatch: pytest.MonkeyPatch,
76 model_setup: tuple[str, str, str, int, float],
77 ):
78 """Test the correctness of EPLB speculative decoding with GSM8K dataset.
79 Applicable to MoE models with mtp or eagle spec decode.
80 """
81 method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup
82
83 TASK = "gsm8k"
84 FILTER = "exact_match,strict-match"
85 RTOL = 0.03
86
87 model_args = get_model_args(
88 model_name=model_name,
89 spec_model_name=spec_model_name,
90 spec_method=method,
91 tp_size=tp_size,
92 model_max_len=4096,
93 )
94
95 results = lm_eval.simple_evaluate(
96 model="vllm",
97 model_args=model_args,
98 tasks=TASK,
99 batch_size=64,
100 num_fewshot=8,
101 )
102 measured_value = results["results"][TASK][FILTER]
103 assert (
104 measured_value - RTOL < expected_gsm8k_value
105 and measured_value + RTOL > expected_gsm8k_value
106 ), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"
107
108
109 @large_gpu_mark(min_gb=80)
110 def test_eplb_spec_decode_qwen3_next_mtp_async() -> None:
111 """Ensure async EPLB works with MTP speculative decoding for Qwen3-Next."""
112 TASK = "gsm8k"
113 FILTER = "exact_match,strict-match"
114 RTOL = 0.03
115 expected_gsm8k_value = 0.86
116
117 model_args = get_model_args(
118 model_name="Qwen/Qwen3-Next-80B-A3B-Instruct",
119 spec_model_name=None,
120 spec_method="mtp",
121 tp_size=4,
122 model_max_len=4096,
123 use_async=True,
124 )
125
126 results = lm_eval.simple_evaluate(
127 model="vllm",
128 model_args=model_args,
129 tasks=TASK,
130 batch_size=64,
131 num_fewshot=8,
132 )
133 measured_value = results["results"][TASK][FILTER]
134 assert (
135 measured_value - RTOL < expected_gsm8k_value
136 and measured_value + RTOL > expected_gsm8k_value
137 ), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"