[Gemma4] Add test for MTP models (#24552)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
b437a0d066
commit
9040feebd8
@@ -0,0 +1,165 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=300, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def get_server_info(base_url: str) -> dict:
|
||||
response = requests.get(base_url + "/server_info", timeout=10)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
|
||||
try:
|
||||
info = get_server_info(base_url)
|
||||
except Exception:
|
||||
return None
|
||||
internal_states = info.get("internal_states") or []
|
||||
if not internal_states:
|
||||
return None
|
||||
value = internal_states[0].get("avg_spec_accept_length")
|
||||
if value is None:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
class TestFrozenKVMTP(CustomTestCase):
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def _server_env(cls) -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
env["SGLANG_ENABLE_SPEC_V2"] = "0"
|
||||
env["SGLANG_ENABLE_ASYNC_ASSERT"] = "1"
|
||||
return env
|
||||
|
||||
@classmethod
|
||||
def _common_server_args(cls) -> list[str]:
|
||||
args = [
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--mem-fraction-static",
|
||||
"0.55",
|
||||
"--max-running-requests",
|
||||
"16",
|
||||
"--context-length",
|
||||
"2048",
|
||||
"--max-total-tokens",
|
||||
"32768",
|
||||
"--skip-server-warmup",
|
||||
]
|
||||
return args
|
||||
|
||||
@classmethod
|
||||
def _server_args(cls, topk: int) -> list[str]:
|
||||
return [
|
||||
"--speculative-algorithm",
|
||||
"NEXTN",
|
||||
"--speculative-draft-model-path",
|
||||
"google/gemma-4-E4B-it-assistant",
|
||||
"--speculative-num-steps",
|
||||
"5",
|
||||
"--speculative-eagle-topk",
|
||||
str(topk),
|
||||
"--speculative-num-draft-tokens",
|
||||
str({1: 6, 3: 12}[topk]),
|
||||
] + cls._common_server_args()
|
||||
|
||||
@classmethod
|
||||
def _gsm8k_args(cls) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
base_url=cls.base_url,
|
||||
model="google/gemma-4-E4B-it",
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=200,
|
||||
num_threads=128,
|
||||
num_shots=5,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _stop_process(process) -> None:
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _run_gsm8k_mtp(self, topk: int) -> None:
|
||||
process = None
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
"google/gemma-4-E4B-it",
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
|
||||
env=self._server_env(),
|
||||
other_args=self._server_args(topk),
|
||||
)
|
||||
requests.get(self.base_url + "/flush_cache", timeout=30)
|
||||
|
||||
server_info = get_server_info(self.base_url)
|
||||
self.assertEqual(
|
||||
server_info.get("speculative_eagle_topk"),
|
||||
topk,
|
||||
f"E4B: server did not start with topk={topk}",
|
||||
)
|
||||
self.assertFalse(
|
||||
bool(server_info.get("disable_cuda_graph")),
|
||||
f"E4B/topk{topk}: CUDA graph is disabled",
|
||||
)
|
||||
|
||||
metrics = run_eval(self._gsm8k_args())
|
||||
mtp_score = float(metrics["score"])
|
||||
avg_accept = get_avg_spec_accept_length(self.base_url)
|
||||
finally:
|
||||
if process is not None:
|
||||
self._stop_process(process)
|
||||
|
||||
print(
|
||||
f"[Frozen-KV MTP E4B topk={topk}] "
|
||||
f"score={mtp_score:.4f} threshold={0.65:.4f} "
|
||||
f"avg_spec_accept_length={avg_accept}"
|
||||
)
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### Frozen-KV MTP E4B topk={topk}\n"
|
||||
f"score={mtp_score:.4f}\n"
|
||||
f"threshold={0.65:.4f}\n"
|
||||
f"avg_spec_accept_length={avg_accept}\n"
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(mtp_score, 0.65)
|
||||
self.assertIsNotNone(avg_accept)
|
||||
self.assertGreaterEqual(
|
||||
avg_accept,
|
||||
1.5,
|
||||
f"E4B/topk{topk}: accept length too low",
|
||||
)
|
||||
|
||||
def test_gsm8k_mtp(self) -> None:
|
||||
for topk in (1, 3):
|
||||
with self.subTest(topk=topk):
|
||||
self._run_gsm8k_mtp(topk)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,185 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=720, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
MODEL_NAME = "26B-A4B"
|
||||
TARGET_PATH = "google/gemma-4-26B-A4B-it"
|
||||
ASSISTANT_PATH = "google/gemma-4-26B-A4B-it-assistant"
|
||||
TENSOR_PARALLEL_SIZE = 2
|
||||
|
||||
TOPKS = (1, 3)
|
||||
DRAFT_TOKENS_BY_TOPK = {1: 6, 3: 12}
|
||||
GSM8K_NUM_EXAMPLES = 200
|
||||
GSM8K_NUM_THREADS = 128
|
||||
GSM8K_SCORE_MARGIN = 0.03
|
||||
SERVER_LAUNCH_TIMEOUT = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3
|
||||
|
||||
# Initial values are seeded from current Gemma4 GSM8K observations in the
|
||||
# cookbook. Replace each top-k entry with exact MTP first-200-sample scores as
|
||||
# CI calibration data becomes available.
|
||||
OBSERVED_GSM8K_SCORES = {1: 0.450, 3: 0.450}
|
||||
GSM8K_SCORE_THRESHOLD = min(OBSERVED_GSM8K_SCORES.values()) - GSM8K_SCORE_MARGIN
|
||||
ACCEPT_LENGTH_THRESHOLD = 1.5
|
||||
|
||||
|
||||
def get_server_info(base_url: str) -> dict:
|
||||
response = requests.get(base_url + "/server_info", timeout=10)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
|
||||
try:
|
||||
info = get_server_info(base_url)
|
||||
except Exception:
|
||||
return None
|
||||
internal_states = info.get("internal_states") or []
|
||||
if not internal_states:
|
||||
return None
|
||||
value = internal_states[0].get("avg_spec_accept_length")
|
||||
if value is None:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
class TestGemma4MTP26BA4B(CustomTestCase):
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def _server_env(cls) -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
env["SGLANG_ENABLE_SPEC_V2"] = "0"
|
||||
return env
|
||||
|
||||
@classmethod
|
||||
def _common_server_args(cls) -> list[str]:
|
||||
args = [
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--mem-fraction-static",
|
||||
"0.55",
|
||||
"--max-running-requests",
|
||||
"16",
|
||||
"--context-length",
|
||||
"2048",
|
||||
"--max-total-tokens",
|
||||
"32768",
|
||||
"--skip-server-warmup",
|
||||
]
|
||||
if TENSOR_PARALLEL_SIZE > 1:
|
||||
args += ["--tp-size", str(TENSOR_PARALLEL_SIZE)]
|
||||
return args
|
||||
|
||||
@classmethod
|
||||
def _server_args(cls, topk: int) -> list[str]:
|
||||
return [
|
||||
"--speculative-algorithm",
|
||||
"NEXTN",
|
||||
"--speculative-draft-model-path",
|
||||
ASSISTANT_PATH,
|
||||
"--speculative-num-steps",
|
||||
"5",
|
||||
"--speculative-eagle-topk",
|
||||
str(topk),
|
||||
"--speculative-num-draft-tokens",
|
||||
str(DRAFT_TOKENS_BY_TOPK[topk]),
|
||||
] + cls._common_server_args()
|
||||
|
||||
@classmethod
|
||||
def _gsm8k_args(cls) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
base_url=cls.base_url,
|
||||
model=TARGET_PATH,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=GSM8K_NUM_EXAMPLES,
|
||||
num_threads=GSM8K_NUM_THREADS,
|
||||
num_shots=5,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _stop_process(process) -> None:
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _run_gsm8k_mtp(self, topk: int) -> None:
|
||||
process = None
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
TARGET_PATH,
|
||||
self.base_url,
|
||||
timeout=SERVER_LAUNCH_TIMEOUT,
|
||||
env=self._server_env(),
|
||||
other_args=self._server_args(topk),
|
||||
)
|
||||
requests.get(self.base_url + "/flush_cache", timeout=30)
|
||||
|
||||
server_info = get_server_info(self.base_url)
|
||||
self.assertEqual(
|
||||
server_info.get("speculative_eagle_topk"),
|
||||
topk,
|
||||
f"{MODEL_NAME}: server did not start with topk={topk}",
|
||||
)
|
||||
self.assertFalse(
|
||||
bool(server_info.get("disable_cuda_graph")),
|
||||
f"{MODEL_NAME}/topk{topk}: CUDA graph is disabled",
|
||||
)
|
||||
|
||||
metrics = run_eval(self._gsm8k_args())
|
||||
mtp_score = float(metrics["score"])
|
||||
avg_accept = get_avg_spec_accept_length(self.base_url)
|
||||
finally:
|
||||
if process is not None:
|
||||
self._stop_process(process)
|
||||
|
||||
print(
|
||||
f"[Gemma4 {MODEL_NAME} topk={topk}] "
|
||||
f"score={mtp_score:.4f} threshold={GSM8K_SCORE_THRESHOLD:.4f} "
|
||||
f"avg_spec_accept_length={avg_accept}"
|
||||
)
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### Gemma4 {MODEL_NAME} MTP topk={topk}\n"
|
||||
f"score={mtp_score:.4f}\n"
|
||||
f"threshold={GSM8K_SCORE_THRESHOLD:.4f}\n"
|
||||
f"avg_spec_accept_length={avg_accept}\n"
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(mtp_score, GSM8K_SCORE_THRESHOLD)
|
||||
self.assertIsNotNone(avg_accept)
|
||||
self.assertGreaterEqual(
|
||||
avg_accept,
|
||||
ACCEPT_LENGTH_THRESHOLD,
|
||||
f"{MODEL_NAME}/topk{topk}: accept length too low",
|
||||
)
|
||||
|
||||
def test_gsm8k_mtp(self) -> None:
|
||||
for topk in TOPKS:
|
||||
with self.subTest(topk=topk):
|
||||
self._run_gsm8k_mtp(topk)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,185 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=720, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
MODEL_NAME = "31B"
|
||||
TARGET_PATH = "google/gemma-4-31B-it"
|
||||
ASSISTANT_PATH = "google/gemma-4-31B-it-assistant"
|
||||
TENSOR_PARALLEL_SIZE = 2
|
||||
|
||||
TOPKS = (1, 3)
|
||||
DRAFT_TOKENS_BY_TOPK = {1: 6, 3: 12}
|
||||
GSM8K_NUM_EXAMPLES = 200
|
||||
GSM8K_NUM_THREADS = 128
|
||||
GSM8K_SCORE_MARGIN = 0.03
|
||||
SERVER_LAUNCH_TIMEOUT = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3
|
||||
|
||||
# Initial values are seeded from current Gemma4 GSM8K observations in the
|
||||
# cookbook. Replace each top-k entry with exact MTP first-200-sample scores as
|
||||
# CI calibration data becomes available.
|
||||
OBSERVED_GSM8K_SCORES = {1: 0.805, 3: 0.805}
|
||||
GSM8K_SCORE_THRESHOLD = min(OBSERVED_GSM8K_SCORES.values()) - GSM8K_SCORE_MARGIN
|
||||
ACCEPT_LENGTH_THRESHOLD = 1.5
|
||||
|
||||
|
||||
def get_server_info(base_url: str) -> dict:
|
||||
response = requests.get(base_url + "/server_info", timeout=10)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
|
||||
try:
|
||||
info = get_server_info(base_url)
|
||||
except Exception:
|
||||
return None
|
||||
internal_states = info.get("internal_states") or []
|
||||
if not internal_states:
|
||||
return None
|
||||
value = internal_states[0].get("avg_spec_accept_length")
|
||||
if value is None:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
class TestGemma4MTP31B(CustomTestCase):
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def _server_env(cls) -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
env["SGLANG_ENABLE_SPEC_V2"] = "0"
|
||||
return env
|
||||
|
||||
@classmethod
|
||||
def _common_server_args(cls) -> list[str]:
|
||||
args = [
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--mem-fraction-static",
|
||||
"0.55",
|
||||
"--max-running-requests",
|
||||
"16",
|
||||
"--context-length",
|
||||
"2048",
|
||||
"--max-total-tokens",
|
||||
"32768",
|
||||
"--skip-server-warmup",
|
||||
]
|
||||
if TENSOR_PARALLEL_SIZE > 1:
|
||||
args += ["--tp-size", str(TENSOR_PARALLEL_SIZE)]
|
||||
return args
|
||||
|
||||
@classmethod
|
||||
def _server_args(cls, topk: int) -> list[str]:
|
||||
return [
|
||||
"--speculative-algorithm",
|
||||
"NEXTN",
|
||||
"--speculative-draft-model-path",
|
||||
ASSISTANT_PATH,
|
||||
"--speculative-num-steps",
|
||||
"5",
|
||||
"--speculative-eagle-topk",
|
||||
str(topk),
|
||||
"--speculative-num-draft-tokens",
|
||||
str(DRAFT_TOKENS_BY_TOPK[topk]),
|
||||
] + cls._common_server_args()
|
||||
|
||||
@classmethod
|
||||
def _gsm8k_args(cls) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
base_url=cls.base_url,
|
||||
model=TARGET_PATH,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=GSM8K_NUM_EXAMPLES,
|
||||
num_threads=GSM8K_NUM_THREADS,
|
||||
num_shots=5,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _stop_process(process) -> None:
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _run_gsm8k_mtp(self, topk: int) -> None:
|
||||
process = None
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
TARGET_PATH,
|
||||
self.base_url,
|
||||
timeout=SERVER_LAUNCH_TIMEOUT,
|
||||
env=self._server_env(),
|
||||
other_args=self._server_args(topk),
|
||||
)
|
||||
requests.get(self.base_url + "/flush_cache", timeout=30)
|
||||
|
||||
server_info = get_server_info(self.base_url)
|
||||
self.assertEqual(
|
||||
server_info.get("speculative_eagle_topk"),
|
||||
topk,
|
||||
f"{MODEL_NAME}: server did not start with topk={topk}",
|
||||
)
|
||||
self.assertFalse(
|
||||
bool(server_info.get("disable_cuda_graph")),
|
||||
f"{MODEL_NAME}/topk{topk}: CUDA graph is disabled",
|
||||
)
|
||||
|
||||
metrics = run_eval(self._gsm8k_args())
|
||||
mtp_score = float(metrics["score"])
|
||||
avg_accept = get_avg_spec_accept_length(self.base_url)
|
||||
finally:
|
||||
if process is not None:
|
||||
self._stop_process(process)
|
||||
|
||||
print(
|
||||
f"[Gemma4 {MODEL_NAME} topk={topk}] "
|
||||
f"score={mtp_score:.4f} threshold={GSM8K_SCORE_THRESHOLD:.4f} "
|
||||
f"avg_spec_accept_length={avg_accept}"
|
||||
)
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### Gemma4 {MODEL_NAME} MTP topk={topk}\n"
|
||||
f"score={mtp_score:.4f}\n"
|
||||
f"threshold={GSM8K_SCORE_THRESHOLD:.4f}\n"
|
||||
f"avg_spec_accept_length={avg_accept}\n"
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(mtp_score, GSM8K_SCORE_THRESHOLD)
|
||||
self.assertIsNotNone(avg_accept)
|
||||
self.assertGreaterEqual(
|
||||
avg_accept,
|
||||
ACCEPT_LENGTH_THRESHOLD,
|
||||
f"{MODEL_NAME}/topk{topk}: accept length too low",
|
||||
)
|
||||
|
||||
def test_gsm8k_mtp(self) -> None:
|
||||
for topk in TOPKS:
|
||||
with self.subTest(topk=topk):
|
||||
self._run_gsm8k_mtp(topk)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user