feat: support EPD disaggregation (#12263)
Co-authored-by: liusy58 <liusy58@linux.alibaba.com> Co-authored-by: ZhengWG <zwg0606@gmail.com> Co-authored-by: Nicholas <45984215+liusy58@users.noreply.github.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
This commit is contained in:
Tianyu Guo
co-authored by
liusy58
ZhengWG
Nicholas
Shangming Cai
Yuhao Yang
liusy58
ZhengWG
Nicholas
Shangming Cai
Yuhao Yang
parent
a9ce1623cd
commit
9acb21ae27
@@ -152,6 +152,7 @@ suites = {
|
||||
TestFile("test_multi_instance_release_memory_occupation.py", 64),
|
||||
TestFile("test_pp_single_node.py", 800),
|
||||
TestFile("test_piecewise_cuda_graph.py", 1200),
|
||||
TestFile("test_epd_disaggregation.py", 600),
|
||||
],
|
||||
"per-commit-8-gpu-h200": [
|
||||
TestFile("test_deepseek_v3_basic.py", 275),
|
||||
|
||||
@@ -0,0 +1,424 @@
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
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,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
]
|
||||
|
||||
subprocess.run(cmd, check=True, 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,
|
||||
]
|
||||
|
||||
subprocess.run(cmd, check=True, 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()
|
||||
Reference in New Issue
Block a user