Migrate 4-GPU/8-GPU workflow jobs to stage-c and add CI registry decorators (#17299)
This commit is contained in:
Alison Shao
@@ -0,0 +1,96 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
popen_launch_pd_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-gb200")
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeAARCH64Accuracy(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
os.environ["SGLANG_MOONCAKE_CUSTOM_MEM_POOL"] = "true"
|
||||
os.environ["MC_FORCE_MNNVL"] = "true"
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
os.environ.pop("SGLANG_MOONCAKE_CUSTOM_MEM_POOL")
|
||||
os.environ.pop("MC_FORCE_MNNVL")
|
||||
super().tearDownClass()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"2",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"2",
|
||||
"--base-gpu-id",
|
||||
"2",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.62)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,306 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
popen_launch_pd_server,
|
||||
try_cached_model,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=600, suite="stage-c-test-8-gpu-h20")
|
||||
|
||||
|
||||
class TestDisaggregationMooncakePrefillLargerTP(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
# Temporarily disable JIT DeepGEMM
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
|
||||
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"4",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"2",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.60)
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeDecodeLargerTP(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
# Temporarily disable JIT DeepGEMM
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
|
||||
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"2",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"4",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.60)
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeMHAPrefillLargerTP(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
# Temporarily disable JIT DeepGEMM
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
|
||||
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"4",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"2",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.60)
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeMHADecodeLargerTP(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
# Temporarily disable JIT DeepGEMM
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
|
||||
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"2",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"4",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,101 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
popen_launch_pd_server,
|
||||
try_cached_model,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=155, suite="stage-c-test-8-gpu-h20")
|
||||
|
||||
|
||||
class TestDisaggregationDPAttention(PDDisaggregationServerBase):
|
||||
PREFILL_DP_SIZE = 4
|
||||
DECODE_DP_SIZE = 4
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
# Temporarily disable JIT DeepGEMM
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
|
||||
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
str(cls.PREFILL_DP_SIZE),
|
||||
"--dp",
|
||||
str(cls.PREFILL_DP_SIZE),
|
||||
"--enable-dp-attention",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
str(cls.DECODE_DP_SIZE),
|
||||
"--dp",
|
||||
str(cls.DECODE_DP_SIZE),
|
||||
"--enable-dp-attention",
|
||||
"--base-gpu-id",
|
||||
str(cls.PREFILL_DP_SIZE),
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=1400,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,240 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
is_in_ci,
|
||||
popen_launch_pd_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=400, suite="stage-c-test-8-gpu-h200", disabled="TCP fallback flaky"
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "Temporarily disable the flaky test.")
|
||||
class TestDisaggregationHybridAttentionMamba(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"4",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"4",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.93)
|
||||
|
||||
|
||||
class TestDisaggregationHybridAttentionMambaExtraBuffer(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"4",
|
||||
"--mamba-scheduler-strategy",
|
||||
"extra_buffer",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"4",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
"--mamba-scheduler-strategy",
|
||||
"extra_buffer",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.93)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
is_in_ci(),
|
||||
"Temporarily disable the flaky test: tcp fallback is not stable currently.",
|
||||
)
|
||||
class TestDisaggregationHybridAttentionMambaDPDecode(PDDisaggregationServerBase):
|
||||
"""Test with prefill tp=2 and decode tp=2/dp=2 with dp-attention enabled."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"2",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"2",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--enable-dp-lm-head",
|
||||
"--base-gpu-id",
|
||||
"2",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.93)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,243 @@
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
popen_launch_pd_server,
|
||||
try_cached_model,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=180, suite="stage-c-test-8-gpu-h20")
|
||||
|
||||
|
||||
class TestDisaggregationPrefillPPAccuracy(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--pp-size",
|
||||
"2",
|
||||
"--disable-overlap-schedule",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.24)
|
||||
# Wait a little bit so that the memory check happens.
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
class TestDisaggregationPrefillPPDynamicChunkAccuracy(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--pp-size",
|
||||
"2",
|
||||
"--disable-overlap-schedule",
|
||||
"--enable-dynamic-chunking",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.24)
|
||||
# Wait a little bit so that the memory check happens.
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
class TestDisaggregationDecodePPAccuracy(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
|
||||
|
||||
# Non blocking start servers
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--pp-size",
|
||||
"2",
|
||||
"--disable-overlap-schedule",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--pp-size",
|
||||
"2",
|
||||
"--base-gpu-id",
|
||||
"4",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{self.base_host}",
|
||||
port=int(self.lb_port),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.24)
|
||||
# Wait a little bit so that the memory check happens.
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,429 @@
|
||||
import os
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=150, suite="stage-c-test-4-gpu-h100")
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "Skipping in CI to reduce multi-GPU runtime")
|
||||
class TestEPDDisaggregationOneEncoder(PDDisaggregationServerBase):
|
||||
"""Test EPD disaggregation with single encode server"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
|
||||
cls.encode_port = f"{int(cls.lb_port) + 300}"
|
||||
cls.encode_url = f"http://{cls.base_host}:{cls.encode_port}"
|
||||
|
||||
print(
|
||||
f"Setting up EPD (one encoder): encode={cls.encode_port}, "
|
||||
f"prefill={cls.prefill_port}, decode={cls.decode_port}"
|
||||
)
|
||||
|
||||
# Start servers in order: encode -> prefill/decode
|
||||
cls.start_encode()
|
||||
prefill_thread = threading.Thread(target=cls.start_prefill)
|
||||
decode_thread = threading.Thread(target=cls.start_decode)
|
||||
prefill_thread.start()
|
||||
decode_thread.start()
|
||||
prefill_thread.join()
|
||||
decode_thread.join()
|
||||
|
||||
# Wait for all servers to be ready
|
||||
cls.wait_server_ready(cls.encode_url + "/health")
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
# Set OpenAI API key and base URL environment variables. Needed for lmms-eval to work.
|
||||
cls.api_key = "sk-123456"
|
||||
os.environ["OPENAI_API_KEY"] = cls.api_key
|
||||
os.environ["OPENAI_API_BASE"] = f"{cls.lb_url}/v1"
|
||||
|
||||
@classmethod
|
||||
def start_encode(cls):
|
||||
"""Start encode server for multimodal processing"""
|
||||
encode_args = [
|
||||
"--trust-remote-code",
|
||||
"--encoder-only",
|
||||
"--encoder-transfer-backend",
|
||||
"zmq_to_scheduler",
|
||||
"--tp",
|
||||
"1",
|
||||
"--port",
|
||||
cls.encode_port,
|
||||
"--enable-prefix-mm-cache",
|
||||
]
|
||||
cls.process_encode = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.encode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=encode_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
"""Start prefill server with language model only"""
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--language-only",
|
||||
"--encoder-urls",
|
||||
cls.encode_url,
|
||||
"--encoder-transfer-backend",
|
||||
"zmq_to_scheduler",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
"--port",
|
||||
cls.prefill_port,
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
"""Start decode server"""
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"2",
|
||||
"--port",
|
||||
cls.decode_port,
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up all processes"""
|
||||
for process in [
|
||||
cls.process_lb,
|
||||
cls.process_decode,
|
||||
cls.process_prefill,
|
||||
cls.process_encode,
|
||||
]:
|
||||
if process:
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
|
||||
"""
|
||||
Evaluate a VLM on the MMMU validation set with lmms-eval.
|
||||
Reference: test_vlm_models.py
|
||||
|
||||
Args:
|
||||
model_version: Model version/checkpoint to evaluate
|
||||
output_path: Path to save evaluation results
|
||||
limit: Number of samples to evaluate (default: "50" for CI time constraints)
|
||||
"""
|
||||
model = "openai_compatible"
|
||||
tp = 1
|
||||
tasks = "mmmu_val"
|
||||
batch_size = 32
|
||||
log_suffix = "openai_compatible"
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
model_args = f'model_version="{model_version}",' f"tp={tp}"
|
||||
|
||||
cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"lmms_eval",
|
||||
"--model",
|
||||
model,
|
||||
"--model_args",
|
||||
model_args,
|
||||
"--tasks",
|
||||
tasks,
|
||||
"--batch_size",
|
||||
str(batch_size),
|
||||
"--log_samples",
|
||||
"--log_samples_suffix",
|
||||
log_suffix,
|
||||
"--output_path",
|
||||
str(output_path),
|
||||
"--limit",
|
||||
limit,
|
||||
]
|
||||
|
||||
_run_lmms_eval_with_retry(cmd, timeout=3600)
|
||||
|
||||
def test_mmmu(self):
|
||||
"""Test MMMU evaluation with EPD disaggregation"""
|
||||
import glob
|
||||
import json
|
||||
|
||||
output_path = "./logs/epd_one_encoder_mmmu"
|
||||
self.run_mmmu_eval(self.model, output_path)
|
||||
|
||||
# Get the result file
|
||||
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
||||
if not result_files:
|
||||
result_files = glob.glob(f"{output_path}/*.json")
|
||||
|
||||
if not result_files:
|
||||
self.fail(f"No JSON result files found in {output_path}")
|
||||
|
||||
result_file_path = result_files[0]
|
||||
with open(result_file_path, "r") as f:
|
||||
result = json.load(f)
|
||||
print(f"MMMU result: {result}")
|
||||
|
||||
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
||||
print(f"MMMU accuracy: {mmmu_accuracy:.4f}")
|
||||
|
||||
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
|
||||
self.assertGreater(mmmu_accuracy, 0.40)
|
||||
|
||||
|
||||
class TestEPDDisaggregationMultiEncoders(PDDisaggregationServerBase):
|
||||
"""
|
||||
Test EPD disaggregation with multiple encode servers for load balancing.
|
||||
Both encode servers run on GPU 0 (different ports) for testing load distribution.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
|
||||
cls.encode_port1 = f"{int(cls.lb_port) + 300}"
|
||||
cls.encode_port2 = f"{int(cls.lb_port) + 301}"
|
||||
cls.encode_url1 = f"http://{cls.base_host}:{cls.encode_port1}"
|
||||
cls.encode_url2 = f"http://{cls.base_host}:{cls.encode_port2}"
|
||||
|
||||
print(
|
||||
f"Setting up EPD (multiple encoders): encode1={cls.encode_port1}, "
|
||||
f"encode2={cls.encode_port2}, prefill={cls.prefill_port}, decode={cls.decode_port}"
|
||||
)
|
||||
|
||||
# Start two encode servers on GPU 0/1
|
||||
encode1_thread = threading.Thread(
|
||||
target=cls.start_encode_server, args=(cls.encode_port1, 0)
|
||||
)
|
||||
encode2_thread = threading.Thread(
|
||||
target=cls.start_encode_server, args=(cls.encode_port2, 1)
|
||||
)
|
||||
encode1_thread.start()
|
||||
encode2_thread.start()
|
||||
encode1_thread.join()
|
||||
encode2_thread.join()
|
||||
|
||||
prefill_thread = threading.Thread(target=cls.start_prefill)
|
||||
decode_thread = threading.Thread(target=cls.start_decode)
|
||||
prefill_thread.start()
|
||||
decode_thread.start()
|
||||
prefill_thread.join()
|
||||
decode_thread.join()
|
||||
|
||||
cls.wait_server_ready(cls.encode_url1 + "/health")
|
||||
cls.wait_server_ready(cls.encode_url2 + "/health")
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
# Set OpenAI API key and base URL environment variables. Needed for lmms-eval to work.
|
||||
cls.api_key = "sk-123456"
|
||||
os.environ["OPENAI_API_KEY"] = cls.api_key
|
||||
os.environ["OPENAI_API_BASE"] = f"{cls.lb_url}/v1"
|
||||
|
||||
@classmethod
|
||||
def start_encode_server(cls, port, gpu_id):
|
||||
"""Start an encode server on specific port and GPU"""
|
||||
encode_args = [
|
||||
"--trust-remote-code",
|
||||
"--encoder-only",
|
||||
"--encoder-transfer-backend",
|
||||
"zmq_to_scheduler",
|
||||
"--tp",
|
||||
"1",
|
||||
"--port",
|
||||
port,
|
||||
"--enable-prefix-mm-cache",
|
||||
]
|
||||
# Only set base-gpu-id if not using GPU 0
|
||||
if gpu_id != 0:
|
||||
encode_args.extend(["--base-gpu-id", str(gpu_id)])
|
||||
|
||||
process = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=f"http://{cls.base_host}:{port}",
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=encode_args,
|
||||
)
|
||||
if port == cls.encode_port1:
|
||||
cls.process_encode1 = process
|
||||
else:
|
||||
cls.process_encode2 = process
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
"""Start prefill server with multiple encode URLs"""
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--language-only",
|
||||
"--encoder-urls",
|
||||
cls.encode_url1,
|
||||
cls.encode_url2,
|
||||
"--encoder-transfer-backend",
|
||||
"zmq_to_scheduler",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"2",
|
||||
"--port",
|
||||
cls.prefill_port,
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_prefill = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
"""Start decode server"""
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"3",
|
||||
"--port",
|
||||
cls.decode_port,
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
cls.process_decode = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up all processes"""
|
||||
for process in [
|
||||
cls.process_lb,
|
||||
cls.process_decode,
|
||||
cls.process_prefill,
|
||||
cls.process_encode1,
|
||||
cls.process_encode2,
|
||||
]:
|
||||
if process:
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
|
||||
"""
|
||||
Evaluate a VLM on the MMMU validation set with lmms-eval.
|
||||
Reference: test_vlm_models.py
|
||||
|
||||
Args:
|
||||
model_version: Model version/checkpoint to evaluate
|
||||
output_path: Path to save evaluation results
|
||||
limit: Number of samples to evaluate (default: "50" for CI time constraints)
|
||||
"""
|
||||
model = "openai_compatible"
|
||||
tp = 1
|
||||
tasks = "mmmu_val"
|
||||
batch_size = 32
|
||||
log_suffix = "openai_compatible"
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
model_args = f'model_version="{model_version}",' f"tp={tp}"
|
||||
|
||||
cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"lmms_eval",
|
||||
"--model",
|
||||
model,
|
||||
"--model_args",
|
||||
model_args,
|
||||
"--tasks",
|
||||
tasks,
|
||||
"--batch_size",
|
||||
str(batch_size),
|
||||
"--log_samples",
|
||||
"--log_samples_suffix",
|
||||
log_suffix,
|
||||
"--output_path",
|
||||
str(output_path),
|
||||
"--limit",
|
||||
limit,
|
||||
]
|
||||
|
||||
_run_lmms_eval_with_retry(cmd, timeout=3600)
|
||||
|
||||
def test_mmmu(self):
|
||||
"""Test MMMU evaluation with EPD disaggregation (multiple encoders)"""
|
||||
import glob
|
||||
import json
|
||||
|
||||
output_path = "./logs/epd_multi_encoder_mmmu"
|
||||
self.run_mmmu_eval(self.model, output_path)
|
||||
|
||||
# Get the result file
|
||||
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
||||
if not result_files:
|
||||
result_files = glob.glob(f"{output_path}/*.json")
|
||||
|
||||
if not result_files:
|
||||
self.fail(f"No JSON result files found in {output_path}")
|
||||
|
||||
result_file_path = result_files[0]
|
||||
with open(result_file_path, "r") as f:
|
||||
result = json.load(f)
|
||||
print(f"MMMU result (multi encoder): {result}")
|
||||
|
||||
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
||||
print(f"MMMU accuracy (multi encoder): {mmmu_accuracy:.4f}")
|
||||
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
|
||||
self.assertGreater(mmmu_accuracy, 0.40)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,427 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_pp_single_node.TestPPAccuracy.test_gsm8k
|
||||
python3 -m unittest test_pp_single_node.TestQwenPPAccuracy.test_pp_consistency
|
||||
python3 -m unittest test_pp_single_node.TestFixedBugs.test_chunked_prefill_with_small_bs
|
||||
python3 -m unittest test_pp_single_node.TestQwenVLPPAccuracy.test_mmmu
|
||||
"""
|
||||
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.bench_one_batch_server import BenchArgs as OneBatchBenchArgs
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_VL_PP,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
run_bench_one_batch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=500, suite="stage-c-test-4-gpu-h100")
|
||||
|
||||
|
||||
class TestPPAccuracy(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = "http://127.0.0.1:23333"
|
||||
cls.process = popen_launch_server(
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp-size",
|
||||
2,
|
||||
"--pp-size",
|
||||
2,
|
||||
"--chunked-prefill-size",
|
||||
256,
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.74)
|
||||
# Wait a little bit so that the memory check happens.
|
||||
time.sleep(4)
|
||||
|
||||
def test_logprob(self):
|
||||
response = requests.post(
|
||||
f"{self.base_url}/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 16,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 5,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
input_token_logprobs = response_json["meta_info"]["input_token_logprobs"]
|
||||
output_token_logprobs = response_json["meta_info"]["output_token_logprobs"]
|
||||
output_top_logprobs = response_json["meta_info"]["output_top_logprobs"]
|
||||
|
||||
assert len(input_token_logprobs) == 6
|
||||
assert len(output_token_logprobs) == 16
|
||||
assert len(output_top_logprobs) == 16
|
||||
|
||||
|
||||
class TestDPAttentionDP2PP2(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"2",
|
||||
"--pp-size",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--dp",
|
||||
"2",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mgsm_en(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mgsm_en",
|
||||
num_examples=None,
|
||||
num_threads=1024,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.8)
|
||||
|
||||
|
||||
class TestQwenVLPPAccuracy(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_VL_PP
|
||||
cls.base_url = "http://127.0.0.1:23333"
|
||||
cls.process = popen_launch_server(
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_VL_PP,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp-size",
|
||||
1,
|
||||
"--pp-size",
|
||||
4,
|
||||
"--chunked-prefill-size",
|
||||
8192,
|
||||
"--enable-multimodal",
|
||||
],
|
||||
)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.65)
|
||||
# Wait a little bit so that the memory check happens.
|
||||
time.sleep(4)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
def test_mmmu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmmu",
|
||||
num_examples=None,
|
||||
num_threads=32,
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.26)
|
||||
|
||||
|
||||
class TestQwenPPAccuracy(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = "http://127.0.0.1:23334" # different ports to avoid conflicts
|
||||
cls.model_name = "Qwen/Qwen3-8B" # replace with your Qwen Model if needed
|
||||
|
||||
def run_gsm8k_test(self, pp_size):
|
||||
process = popen_launch_server(
|
||||
self.model_name,
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--pp-size",
|
||||
pp_size,
|
||||
"--chunked-prefill-size",
|
||||
256,
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
time.sleep(5)
|
||||
return metrics
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
def test_pp_consistency(self):
|
||||
baseline = self.run_gsm8k_test(pp_size=1)
|
||||
pp_metrics = self.run_gsm8k_test(pp_size=2)
|
||||
|
||||
print(f"[Qwen PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
|
||||
|
||||
self.assertGreaterEqual(baseline["accuracy"], 0.74)
|
||||
self.assertGreaterEqual(
|
||||
pp_metrics["accuracy"],
|
||||
baseline["accuracy"] - 0.02,
|
||||
msg=(
|
||||
f"PP accuracy dropped more than 1% compared to baseline. "
|
||||
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestQwenPPTieWeightsAccuracy(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = "http://127.0.0.1:23335" # different ports to avoid conflicts
|
||||
cls.model_name = (
|
||||
"Qwen/Qwen3-0.6B" # qwen3 < 8B all have tie_word_embeddings = True
|
||||
)
|
||||
|
||||
def run_gsm8k_test(self, pp_size):
|
||||
process = popen_launch_server(
|
||||
self.model_name,
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--pp-size",
|
||||
pp_size,
|
||||
"--chunked-prefill-size",
|
||||
256,
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
time.sleep(5)
|
||||
return metrics
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
def test_pp_consistency(self):
|
||||
baseline = self.run_gsm8k_test(pp_size=1)
|
||||
pp_metrics = self.run_gsm8k_test(pp_size=2)
|
||||
|
||||
print(f"[Qwen PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
|
||||
|
||||
self.assertGreaterEqual(baseline["accuracy"], 0.38)
|
||||
self.assertGreaterEqual(
|
||||
pp_metrics["accuracy"],
|
||||
baseline["accuracy"] - 0.02,
|
||||
msg=(
|
||||
f"PP accuracy dropped more than 1% compared to baseline. "
|
||||
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestQwenMoePPAccuracy(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = "http://127.0.0.1:23336" # different ports to avoid conflicts
|
||||
cls.model_name = "Qwen/Qwen3-30B-A3B" # replace with your Qwen Model if needed
|
||||
|
||||
def run_gsm8k_test(self, pp_size):
|
||||
process = popen_launch_server(
|
||||
self.model_name,
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--pp-size",
|
||||
pp_size,
|
||||
"--chunked-prefill-size",
|
||||
256,
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
time.sleep(5)
|
||||
return metrics
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
def test_pp_consistency(self):
|
||||
baseline = self.run_gsm8k_test(pp_size=1)
|
||||
pp_metrics = self.run_gsm8k_test(pp_size=2)
|
||||
|
||||
print(f"[Qwen PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
|
||||
|
||||
self.assertGreaterEqual(baseline["accuracy"], 0.74)
|
||||
self.assertGreaterEqual(
|
||||
pp_metrics["accuracy"],
|
||||
baseline["accuracy"] - 0.02,
|
||||
msg=(
|
||||
f"PP accuracy dropped more than 1% compared to baseline. "
|
||||
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestFixedBugs(unittest.TestCase):
|
||||
def test_chunked_prefill_with_small_bs(self):
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
server_args = ServerArgs(model_path=model)
|
||||
bench_args = OneBatchBenchArgs(
|
||||
batch_size=(1,),
|
||||
input_len=(1,),
|
||||
output_len=(1,),
|
||||
base_url=DEFAULT_URL_FOR_TEST,
|
||||
)
|
||||
other_server_args = [
|
||||
"--tp-size",
|
||||
2,
|
||||
"--pp-size",
|
||||
2,
|
||||
"--chunked-prefill",
|
||||
256,
|
||||
"--max-running-requests",
|
||||
2,
|
||||
]
|
||||
run_bench_one_batch_server(
|
||||
model,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
server_args,
|
||||
bench_args,
|
||||
other_server_args,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
is_in_ci(), "Skipping GLM41V PP accuracy test before it gets more stable"
|
||||
)
|
||||
class TestGLM41VPPAccuracy(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp-size",
|
||||
1,
|
||||
"--pp-size",
|
||||
2,
|
||||
"--chunked-prefill-size",
|
||||
8192,
|
||||
"--enable-multimodal",
|
||||
"--reasoning-parser",
|
||||
"glm45",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmmu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmmu",
|
||||
num_examples=None,
|
||||
num_threads=32,
|
||||
response_answer_regex="<\|begin_of_box\|>(.*)<\|end_of_box\|>",
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.45)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user