ci: migrate 2-GPU tests to test/registered/ (#16529)
This commit is contained in:
Alison Shao
@@ -1,88 +0,0 @@
|
||||
"""
|
||||
Benchmark tests for HiCache Storage with 3FS backend.
|
||||
Usage:
|
||||
python3 -m pytest test/srt/hicache/test_hicache_storage_3fs_backend.py -v
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class HiCacheStorage3FSBackendBaseMixin(HiCacheStorageBaseMixin):
|
||||
"""Base mixin class with common setup and utilities"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
# Create a temporary JSON config file for HF3FS
|
||||
hf3fs_config = {
|
||||
"file_path_prefix": os.path.join(cls.temp_dir, "hicache"),
|
||||
"file_size": 1024 * 1024 * 1024 * 2,
|
||||
"numjobs": 2,
|
||||
"entries": 8,
|
||||
"use_mock_hf3fs_client": True,
|
||||
"hicache_storage_pass_prefix_keys": True,
|
||||
}
|
||||
|
||||
# Write config to temporary file
|
||||
config_file = os.path.join(cls.temp_dir, "hf3fs_config.json")
|
||||
with open(config_file, "w") as f:
|
||||
json.dump(hf3fs_config, f, indent=2)
|
||||
|
||||
server_args = {
|
||||
"--tp-size": 1,
|
||||
"--hicache-ratio": 1.2,
|
||||
"--hicache-storage-backend": "hf3fs",
|
||||
"--hicache-storage-backend-extra-config": json.dumps(hf3fs_config),
|
||||
}
|
||||
|
||||
# Set the environment variable to point to our config file
|
||||
env_vars = {
|
||||
"SGLANG_HICACHE_HF3FS_CONFIG_PATH": config_file,
|
||||
}
|
||||
|
||||
return server_args, env_vars
|
||||
|
||||
|
||||
class TestHf3fsBackendLayerFirstLayout(
|
||||
HiCacheStorage3FSBackendBaseMixin, CustomTestCase
|
||||
):
|
||||
"""Layer first layout tests for HiCache-Hf3fs backend"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args, env_vars = super()._get_additional_server_args_and_env()
|
||||
server_args["--hicache-mem-layout"] = "layer_first"
|
||||
server_args["--hicache-io-backend"] = "direct"
|
||||
server_args["--tp-size"] = 2
|
||||
return server_args, env_vars
|
||||
|
||||
|
||||
class TestHf3fsBackendAccuracy(HiCacheStorage3FSBackendBaseMixin, CustomTestCase):
|
||||
"""Accuracy tests for HiCache-Hf3fs backend"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args, env_vars = super()._get_additional_server_args_and_env()
|
||||
server_args["--hicache-ratio"] = 1.5
|
||||
server_args["--tp-size"] = 2
|
||||
server_args["--hicache-mem-layout"] = "page_first_direct"
|
||||
server_args["--hicache-io-backend"] = "direct"
|
||||
return server_args, env_vars
|
||||
|
||||
def test_eval_accuracy(self):
|
||||
"""Test eval accuracy with cache persistence across cache flushes"""
|
||||
from test_hicache_storage_file_backend import run_eval_accuracy_test
|
||||
|
||||
run_eval_accuracy_test(self)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -1,333 +0,0 @@
|
||||
"""
|
||||
E2E tests for HiCache Storage functionality.
|
||||
Usage:
|
||||
python3 -m pytest test/srt/hicache/test_hicache_storage_e2e.py -v
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from typing import Dict
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.bench_serving import get_tokenizer
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class HiCacheStorageBaseMixin:
|
||||
"""Base mixin class with common setup and utilities"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Set up test environment and launch server once for all tests"""
|
||||
cls.temp_dir = tempfile.mkdtemp()
|
||||
cls.model = cls._get_model_name()
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
parsed_url = urlparse(cls.base_url)
|
||||
cls.base_host = parsed_url.hostname
|
||||
cls.base_port = str(parsed_url.port)
|
||||
|
||||
# Prepare tokenizer for prompt generation
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
# Launch server with HiCache enabled and cache report
|
||||
cls.process = cls._launch_server_with_hicache()
|
||||
cls._wait_for_server_ready()
|
||||
|
||||
print(f"Test server launched successfully at {cls.base_url}")
|
||||
print(f"Cache directory: {cls.temp_dir}")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up test environment"""
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(cls.temp_dir, ignore_errors=True)
|
||||
|
||||
@classmethod
|
||||
def _get_model_name(cls):
|
||||
"""Get model name for the test configuration - override in subclasses"""
|
||||
return DEFAULT_MODEL_NAME_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def _get_base_server_args(cls):
|
||||
"""Get base server arguments - can be extended in subclasses"""
|
||||
extra_config = {
|
||||
"hicache_storage_pass_prefix_keys": True,
|
||||
}
|
||||
return {
|
||||
"--enable-hierarchical-cache": True,
|
||||
"--mem-fraction-static": 0.6,
|
||||
"--hicache-ratio": 1.2,
|
||||
"--page-size": 64,
|
||||
"--enable-cache-report": True,
|
||||
"--hicache-storage-prefetch-policy": "wait_complete",
|
||||
"--hicache-storage-backend": "file",
|
||||
"--hicache-storage-backend-extra-config": json.dumps(extra_config),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
return {}, {"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir}
|
||||
|
||||
@classmethod
|
||||
def _launch_server_with_hicache(cls):
|
||||
"""Launch server with HiCache enabled"""
|
||||
|
||||
additional_server_args, env_vars = cls._get_additional_server_args_and_env()
|
||||
env_vars["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"] = "1"
|
||||
server_args = cls._get_base_server_args()
|
||||
if additional_server_args:
|
||||
server_args.update(additional_server_args)
|
||||
|
||||
final_server_args = []
|
||||
for k, v in server_args.items():
|
||||
if isinstance(v, bool):
|
||||
final_server_args.append(str(k))
|
||||
else:
|
||||
final_server_args.append(str(k))
|
||||
final_server_args.append(str(v))
|
||||
|
||||
print(f"final_server_args: {final_server_args}")
|
||||
|
||||
env_vars = {
|
||||
**os.environ,
|
||||
**env_vars,
|
||||
}
|
||||
|
||||
return popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=final_server_args,
|
||||
env=env_vars,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _wait_for_server_ready(cls, timeout: int = 60) -> bool:
|
||||
"""Wait for server to be ready"""
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < timeout:
|
||||
try:
|
||||
response = requests.get(f"{cls.base_url}/health", timeout=5)
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(2)
|
||||
raise TimeoutError("Server failed to start within timeout")
|
||||
|
||||
def send_request(
|
||||
self, prompt: str, max_tokens: int = 100, temperature: float = 0.0
|
||||
) -> Dict:
|
||||
"""Send a generate request and return response"""
|
||||
response = requests.post(
|
||||
f"{self.base_url}/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": temperature,
|
||||
"max_new_tokens": max_tokens,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
response.status_code,
|
||||
200,
|
||||
f"Request failed: {response.status_code} - {response.text}",
|
||||
)
|
||||
return response.json()
|
||||
|
||||
def get_cached_tokens(self, response_json: Dict) -> int:
|
||||
"""Extract cached tokens count from /generate response"""
|
||||
meta = response_json.get("meta_info", {})
|
||||
return int(meta.get("cached_tokens", 0))
|
||||
|
||||
def flush_cache(self) -> bool:
|
||||
"""Flush device cache to force remote storage access"""
|
||||
try:
|
||||
response = requests.post(f"{self.base_url}/flush_cache", timeout=10)
|
||||
return response.status_code == 200
|
||||
except requests.RequestException:
|
||||
return False
|
||||
|
||||
def gen_prompt(self, token_num: int) -> str:
|
||||
"""Generate a random prompt of specified token length using tokenizer vocabulary."""
|
||||
all_available_tokens = list(self.tokenizer.get_vocab().values())
|
||||
selected_tokens = random.choices(all_available_tokens, k=token_num)
|
||||
return self.tokenizer.decode(selected_tokens)
|
||||
|
||||
def trigger_offloading_and_flush(self):
|
||||
"""Helper method to trigger offloading and flush cache"""
|
||||
# Trigger offloading
|
||||
self.send_request(self.gen_prompt(1), max_tokens=150)
|
||||
|
||||
# Flush device cache to force remote storage access
|
||||
time.sleep(2)
|
||||
self.assertTrue(self.flush_cache(), "Cache flush should succeed")
|
||||
|
||||
def test_basic_backup_and_prefetch(self):
|
||||
"""Test storage and retrieval of large context through remote cache"""
|
||||
print("\n=== Testing Large Context Cache Storage & Retrieval ===")
|
||||
|
||||
# Generate substantial context that will be cached
|
||||
base_prompt = self.gen_prompt(768)
|
||||
|
||||
# First request - populate cache
|
||||
print("Step 1: Populating cache with large context...")
|
||||
response1 = self.send_request(base_prompt, max_tokens=150)
|
||||
self.assertIsNotNone(response1)
|
||||
|
||||
# Flush device cache to force remote storage access
|
||||
self.trigger_offloading_and_flush()
|
||||
|
||||
# Second request with extended prompt - should hit remote cache
|
||||
print("Step 2: Testing cache hit from remote storage...")
|
||||
|
||||
start_time = time.time()
|
||||
response2 = self.send_request(base_prompt, max_tokens=150)
|
||||
retrieval_time = time.time() - start_time
|
||||
|
||||
cached_tokens = self.get_cached_tokens(response2)
|
||||
print(
|
||||
f"Remote cache retrieval time: {retrieval_time:.3f}s, cached_tokens={cached_tokens}"
|
||||
)
|
||||
|
||||
# Assert cached tokens indicate a remote hit
|
||||
self.assertGreater(
|
||||
cached_tokens, 700, "Expected significant cached tokens for remote hit"
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
class TestHiCacheStoragePageFirstLayout(HiCacheStorageBaseMixin, CustomTestCase):
|
||||
"""Page first layout tests for HiCache Storage functionality"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args = {"--hicache-mem-layout": "page_first"}
|
||||
return server_args, {}
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
class TestHiCacheStorageMLA(HiCacheStorageBaseMixin, CustomTestCase):
|
||||
"""MLA Model tests for HiCache Storage functionality"""
|
||||
|
||||
@classmethod
|
||||
def _get_model_name(cls):
|
||||
"""Use MLA model for testing"""
|
||||
return DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args = {"--tp-size": 2}
|
||||
return server_args, {}
|
||||
|
||||
|
||||
class TestHiCacheStoragePageFirstDirectIO(HiCacheStorageBaseMixin, CustomTestCase):
|
||||
"""Page first direct tests for HiCache Storage functionality"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args = {
|
||||
"--hicache-mem-layout": "page_first_direct",
|
||||
"--hicache-io-backend": "direct",
|
||||
"--tp-size": 2,
|
||||
}
|
||||
return server_args, {}
|
||||
|
||||
|
||||
class TestHiCacheStorageAccuracy(HiCacheStorageBaseMixin, CustomTestCase):
|
||||
"""Accuracy tests for HiCache Storage functionality"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args = {
|
||||
"--tp-size": 2,
|
||||
"--hicache-ratio": 1.5,
|
||||
}
|
||||
|
||||
return server_args, {}
|
||||
|
||||
def test_eval_accuracy(self):
|
||||
"""Test eval accuracy with cache persistence across cache flushes"""
|
||||
run_eval_accuracy_test(self)
|
||||
|
||||
|
||||
def run_eval_accuracy_test(test_instance, accuracy_threshold: float = 0.03):
|
||||
"""Generic eval accuracy test with configurable accuracy threshold
|
||||
|
||||
Args:
|
||||
test_instance: The test class instance that provides base_host, base_port, flush_cache, and assert methods
|
||||
"""
|
||||
print("\n=== Testing Eval Accuracy with Cache Persistence ===")
|
||||
|
||||
# First evaluation - populate cache
|
||||
print("Phase 1: Running initial GSM8K evaluation to populate cache...")
|
||||
args_initial = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=50,
|
||||
max_new_tokens=512,
|
||||
parallel=10,
|
||||
host=f"http://{test_instance.base_host}",
|
||||
port=int(test_instance.base_port),
|
||||
)
|
||||
metrics_initial = run_eval_few_shot_gsm8k(args_initial)
|
||||
|
||||
# Flush cache to force remote storage access
|
||||
print("Phase 2: Flushing device cache...")
|
||||
test_instance.assertTrue(test_instance.flush_cache(), "Cache flush should succeed")
|
||||
time.sleep(2)
|
||||
|
||||
# Second evaluation - should use remote cache
|
||||
print("Phase 3: Running second GSM8K evaluation using remote cache...")
|
||||
metrics_cached = run_eval_few_shot_gsm8k(args_initial)
|
||||
|
||||
# Verify accuracy consistency
|
||||
accuracy_diff = abs(metrics_initial["accuracy"] - metrics_cached["accuracy"])
|
||||
print(f"Accuracy difference: {accuracy_diff:.4f}")
|
||||
|
||||
# Assertions
|
||||
test_instance.assertGreater(
|
||||
metrics_initial["accuracy"], 0.6, "Initial accuracy should be reasonable"
|
||||
)
|
||||
test_instance.assertGreater(
|
||||
metrics_cached["accuracy"], 0.6, "Cached accuracy should be reasonable"
|
||||
)
|
||||
test_instance.assertLess(
|
||||
accuracy_diff,
|
||||
accuracy_threshold,
|
||||
"Accuracy should be consistent between cache states",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -1,283 +0,0 @@
|
||||
"""
|
||||
Benchmark tests for HiCache Storage with Mooncake backend.
|
||||
Usage:
|
||||
python3.10 -m pytest test/srt/hicache/test_hicache_storage_mooncake_backend.py -v
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
|
||||
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
CustomTestCase,
|
||||
find_available_port,
|
||||
is_in_ci,
|
||||
)
|
||||
|
||||
|
||||
class HiCacheStorageMooncakeBackendBaseMixin(HiCacheStorageBaseMixin):
|
||||
"""Base mixin class with common setup and utilities"""
|
||||
|
||||
# Default port ranges for Mooncake services - can be overridden in subclasses
|
||||
mooncake_master_port_base = 50051
|
||||
mooncake_metadata_port_base = 8080
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Set up test environment and launch Mooncake services before server setup"""
|
||||
# Find available ports for Mooncake services to avoid conflicts
|
||||
cls.mooncake_master_port = find_available_port(
|
||||
HiCacheStorageMooncakeBackendBaseMixin.mooncake_master_port_base
|
||||
)
|
||||
cls.mooncake_metadata_port = find_available_port(
|
||||
HiCacheStorageMooncakeBackendBaseMixin.mooncake_metadata_port_base
|
||||
)
|
||||
|
||||
# Start Mooncake services first
|
||||
cls._start_mooncake_services()
|
||||
|
||||
# Call parent setup
|
||||
super().setUpClass()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up Mooncake services after server teardown"""
|
||||
# Call parent teardown first
|
||||
super().tearDownClass()
|
||||
|
||||
# Stop Mooncake services
|
||||
cls._stop_mooncake_services()
|
||||
|
||||
@classmethod
|
||||
def _start_mooncake_services(cls):
|
||||
"""Start Mooncake metadata and master services with configurable ports and readiness detection"""
|
||||
print("Starting Mooncake services...")
|
||||
print(
|
||||
f"Using master port: {cls.mooncake_master_port}, metadata port: {cls.mooncake_metadata_port}"
|
||||
)
|
||||
|
||||
# Start metadata service with configurable port
|
||||
try:
|
||||
# Start metadata server with port configuration
|
||||
cls.metadata_service_process = subprocess.Popen(
|
||||
[
|
||||
"python3",
|
||||
"-m",
|
||||
"mooncake.http_metadata_server",
|
||||
"--port",
|
||||
str(cls.mooncake_metadata_port),
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
preexec_fn=os.setsid, # Create new process group
|
||||
)
|
||||
print(
|
||||
f"Mooncake metadata service started on port {cls.mooncake_metadata_port}"
|
||||
)
|
||||
except (FileNotFoundError, subprocess.SubprocessError) as e:
|
||||
print(f"Warning: Could not start Mooncake metadata service: {e}")
|
||||
cls.metadata_service_process = None
|
||||
|
||||
# Start master service with configurable port
|
||||
try:
|
||||
# Start master server with port configuration
|
||||
cls.master_service_process = subprocess.Popen(
|
||||
["mooncake_master", "--port", str(cls.mooncake_master_port)],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
preexec_fn=os.setsid, # Create new process group
|
||||
)
|
||||
print(f"Mooncake master service started on port {cls.mooncake_master_port}")
|
||||
except (FileNotFoundError, subprocess.SubprocessError) as e:
|
||||
print(f"Warning: Could not start Mooncake master service: {e}")
|
||||
cls.master_service_process = None
|
||||
|
||||
# Wait for services to be ready instead of fixed sleep
|
||||
cls._wait_for_mooncake_services_ready()
|
||||
|
||||
@classmethod
|
||||
def _wait_for_mooncake_services_ready(cls, timeout: int = 30) -> bool:
|
||||
"""Wait for Mooncake services to be ready by checking their endpoints"""
|
||||
print("Waiting for Mooncake services to be ready...")
|
||||
|
||||
start_time = time.time()
|
||||
services_ready = False
|
||||
|
||||
while time.time() - start_time < timeout:
|
||||
try:
|
||||
# Check metadata service
|
||||
metadata_ready = False
|
||||
if (
|
||||
cls.metadata_service_process
|
||||
and cls.metadata_service_process.poll() is None
|
||||
):
|
||||
try:
|
||||
# Try to connect to the metadata service
|
||||
metadata_url = (
|
||||
f"http://127.0.0.1:{cls.mooncake_metadata_port}/metadata"
|
||||
)
|
||||
response = requests.get(metadata_url, timeout=2)
|
||||
if response.status_code == 200:
|
||||
metadata_ready = True
|
||||
print("Mooncake metadata service is ready")
|
||||
except (requests.RequestException, ConnectionError):
|
||||
# Service might not be fully started yet
|
||||
pass
|
||||
|
||||
# Check master service (if it has a health endpoint)
|
||||
master_ready = False
|
||||
if (
|
||||
cls.master_service_process
|
||||
and cls.master_service_process.poll() is None
|
||||
):
|
||||
# For now, we'll assume master service is ready if process is running
|
||||
# and it's been a few seconds since startup
|
||||
if (
|
||||
time.time() - start_time > 5
|
||||
): # Give master service time to initialize
|
||||
master_ready = True
|
||||
print("Mooncake master service is ready")
|
||||
|
||||
# Both services should be ready
|
||||
if metadata_ready and master_ready:
|
||||
services_ready = True
|
||||
print("All Mooncake services are ready")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error checking service readiness: {e}")
|
||||
|
||||
time.sleep(2)
|
||||
|
||||
if not services_ready:
|
||||
print(
|
||||
"Warning: Mooncake services may not be fully ready, continuing anyway..."
|
||||
)
|
||||
|
||||
return services_ready
|
||||
|
||||
@classmethod
|
||||
def _stop_mooncake_services(cls):
|
||||
"""Stop Mooncake services"""
|
||||
print("Stopping Mooncake services...")
|
||||
|
||||
# Stop metadata service
|
||||
if hasattr(cls, "metadata_service_process") and cls.metadata_service_process:
|
||||
try:
|
||||
os.killpg(os.getpgid(cls.metadata_service_process.pid), 9)
|
||||
cls.metadata_service_process.wait(timeout=5)
|
||||
print("Mooncake metadata service stopped")
|
||||
except (ProcessLookupError, subprocess.TimeoutExpired, OSError) as e:
|
||||
print(f"Warning: Could not stop Mooncake metadata service: {e}")
|
||||
|
||||
# Stop master service
|
||||
if hasattr(cls, "master_service_process") and cls.master_service_process:
|
||||
try:
|
||||
os.killpg(os.getpgid(cls.master_service_process.pid), 9)
|
||||
cls.master_service_process.wait(timeout=5)
|
||||
print("Mooncake master service stopped")
|
||||
except (ProcessLookupError, subprocess.TimeoutExpired, OSError) as e:
|
||||
print(f"Warning: Could not stop Mooncake master service: {e}")
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
|
||||
server_args = {
|
||||
"--tp-size": 2,
|
||||
"--hicache-ratio": 2,
|
||||
"--hicache-storage-backend": "mooncake",
|
||||
}
|
||||
|
||||
# Set the environment variables for Mooncake using dynamic ports
|
||||
env_vars = {
|
||||
"MOONCAKE_MASTER": f"127.0.0.1:{cls.mooncake_master_port}",
|
||||
"MOONCAKE_PROTOCOL": "tcp",
|
||||
"MC_MS_AUTO_DISC": "0",
|
||||
"MOONCAKE_DEVICE": "",
|
||||
"MOONCAKE_TE_META_DATA_SERVER": f"http://127.0.0.1:{cls.mooncake_metadata_port}/metadata",
|
||||
"MOONCAKE_GLOBAL_SEGMENT_SIZE": "4294967296", # 4 GiB
|
||||
}
|
||||
|
||||
return server_args, env_vars
|
||||
|
||||
|
||||
'''
|
||||
# Same as #10131, layer first layout test TODO(mateng): will make it work
|
||||
class TestMooncakeBackendLayerFirstLayout(
|
||||
HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase
|
||||
):
|
||||
"""Layer first layout tests for HiCache-Mooncake backend"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args, env_vars = super()._get_additional_server_args_and_env()
|
||||
server_args["--hicache-mem-layout"] = "layer_first"
|
||||
server_args["--hicache-io-backend"] = "direct"
|
||||
return server_args, env_vars
|
||||
'''
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
class TestMooncakeBackendPageFirstLayout(
|
||||
HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase
|
||||
):
|
||||
"""Page first layout tests for HiCache-Mooncake backend"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args, env_vars = super()._get_additional_server_args_and_env()
|
||||
server_args["--hicache-mem-layout"] = "page_first"
|
||||
return server_args, env_vars
|
||||
|
||||
|
||||
class TestMooncakeBackendMLAModel(
|
||||
HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase
|
||||
):
|
||||
"""MLA Model tests for HiCache-Mooncake backend"""
|
||||
|
||||
@classmethod
|
||||
def _get_model_name(cls):
|
||||
"""Use MLA model for testing"""
|
||||
return DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args, env_vars = super()._get_additional_server_args_and_env()
|
||||
server_args["--hicache-mem-layout"] = "page_first"
|
||||
server_args["--tp-size"] = 2
|
||||
return server_args, env_vars
|
||||
|
||||
|
||||
class TestMooncakeBackendAccuracy(
|
||||
HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase
|
||||
):
|
||||
"""Accuracy tests for HiCache-Mooncake backend"""
|
||||
|
||||
@classmethod
|
||||
def _get_additional_server_args_and_env(cls):
|
||||
"""Get additional server arguments specific to configuration - override in subclasses"""
|
||||
server_args, env_vars = super()._get_additional_server_args_and_env()
|
||||
server_args["--hicache-ratio"] = 1.5
|
||||
server_args["--tp-size"] = 2
|
||||
server_args["--hicache-mem-layout"] = "page_first_direct"
|
||||
server_args["--hicache-io-backend"] = "direct"
|
||||
return server_args, env_vars
|
||||
|
||||
def test_eval_accuracy(self):
|
||||
"""Test eval accuracy with cache persistence across cache flushes"""
|
||||
from test_hicache_storage_file_backend import run_eval_accuracy_test
|
||||
|
||||
run_eval_accuracy_test(self)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -1,46 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestKimiLinear(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
|
||||
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=["--tp-size", "2", "--trust-remote"],
|
||||
)
|
||||
|
||||
@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(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.88)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,96 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import is_blackwell
|
||||
from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
|
||||
|
||||
class TestNvidiaNemotronNanoV2BF16(GSM8KMixin, DefaultServerBase):
|
||||
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
|
||||
gsm8k_accuracy_thres = 0.87
|
||||
other_args = ["--max-mamba-cache-size", "256"]
|
||||
|
||||
|
||||
class TestNvidiaNemotronNanoV2BF16PP(GSM8KMixin, DefaultServerBase):
|
||||
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
|
||||
gsm8k_accuracy_thres = 0.87
|
||||
other_args = ["--max-mamba-cache-size", "256", "--pp-size", "2"]
|
||||
|
||||
|
||||
class TestNvidiaNemotronNanoV2FP8(GSM8KMixin, DefaultServerBase):
|
||||
gsm8k_accuracy_thres = 0.87
|
||||
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2-FP8"
|
||||
other_args = ["--max-mamba-cache-size", "256"]
|
||||
|
||||
|
||||
@unittest.skipIf(not is_blackwell(), "NVFP4 only supported on blackwell")
|
||||
class TestNvidiaNemotronNanoV2NVFP4(GSM8KMixin, DefaultServerBase):
|
||||
gsm8k_accuracy_thres = 0.855
|
||||
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2-NVFP4"
|
||||
other_args = ["--max-mamba-cache-size", "256"]
|
||||
|
||||
|
||||
@unittest.skip(
|
||||
"STANDALONE speculative decoding does not yet support target and draft models "
|
||||
"with different hidden sizes (Nemotron-9B: 4480, Llama-3.2-1B: 2048)"
|
||||
)
|
||||
class TestNvidiaNemotronNanoV2SpeculativeDecoding(GSM8KMixin, DefaultServerBase):
|
||||
gsm8k_accuracy_thres = 0.87
|
||||
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
|
||||
other_args = [
|
||||
"--speculative-algorithm",
|
||||
"STANDALONE",
|
||||
"--speculative-num-steps",
|
||||
"2",
|
||||
"--speculative-eagle-topk",
|
||||
"3",
|
||||
"--speculative-num-draft-tokens",
|
||||
"5",
|
||||
"--speculative-draft-model-path",
|
||||
"meta-llama/Llama-3.2-1B",
|
||||
"--speculative-draft-load-format",
|
||||
"dummy",
|
||||
"--max-running-requests",
|
||||
"8",
|
||||
"--max-total-tokens",
|
||||
"2048",
|
||||
"--json-model-override-args",
|
||||
'{"vocab_size": 131072}',
|
||||
]
|
||||
|
||||
|
||||
@unittest.skip(
|
||||
"STANDALONE speculative decoding does not yet support target and draft models "
|
||||
"with different hidden sizes (Nemotron-9B: 4480, Llama-3.2-1B: 2048)"
|
||||
)
|
||||
class TestNvidiaNemotronNanoV2SpeculativeDecodingBF16Cache(
|
||||
GSM8KMixin, DefaultServerBase
|
||||
):
|
||||
gsm8k_accuracy_thres = 0.87
|
||||
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
|
||||
other_args = [
|
||||
"--speculative-algorithm",
|
||||
"STANDALONE",
|
||||
"--speculative-num-steps",
|
||||
"2",
|
||||
"--speculative-eagle-topk",
|
||||
"3",
|
||||
"--speculative-num-draft-tokens",
|
||||
"5",
|
||||
"--speculative-draft-model-path",
|
||||
"meta-llama/Llama-3.2-1B",
|
||||
"--speculative-draft-load-format",
|
||||
"dummy",
|
||||
"--max-running-requests",
|
||||
"8",
|
||||
"--max-total-tokens",
|
||||
"2048",
|
||||
"--json-model-override-args",
|
||||
'{"vocab_size": 131072}',
|
||||
"--mamba-ssm-dtype",
|
||||
"bfloat16",
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -16,17 +16,6 @@ suites = {
|
||||
TestFile("test_video_utils.py", 5),
|
||||
TestFile("test_modelopt_export.py", 9),
|
||||
],
|
||||
"per-commit-2-gpu": [
|
||||
TestFile("hicache/test_hicache_storage_3fs_backend.py", 200),
|
||||
TestFile("hicache/test_hicache_storage_file_backend.py", 200),
|
||||
TestFile("hicache/test_hicache_storage_mooncake_backend.py", 300),
|
||||
TestFile("models/test_kimi_linear_models.py", 90),
|
||||
TestFile("models/test_nvidia_nemotron_nano_v2.py", 132),
|
||||
TestFile("test_data_parallelism.py", 73),
|
||||
TestFile("test_disaggregation_basic.py", 400),
|
||||
TestFile("test_dp_attention.py", 350),
|
||||
TestFile("test_load_weights_from_remote_instance.py", 72),
|
||||
],
|
||||
"per-commit-4-gpu": [
|
||||
TestFile("models/test_qwen3_next_models.py", 650),
|
||||
TestFile("test_gpt_oss_4gpu.py", 300),
|
||||
@@ -116,11 +105,6 @@ suite_amd = {
|
||||
# TestFile("test_vision_chunked_prefill.py", 175), # Disabled temporarily and track in #7701
|
||||
# TestFile("test_wave_attention_backend.py", 150), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
|
||||
],
|
||||
"per-commit-amd-mi35x": [],
|
||||
"per-commit-2-gpu-amd": [
|
||||
TestFile("test_data_parallelism.py", 73),
|
||||
TestFile("test_load_weights_from_remote_instance.py", 72),
|
||||
],
|
||||
"per-commit-4-gpu-amd": [
|
||||
TestFile("test_pp_single_node.py", 150),
|
||||
],
|
||||
|
||||
@@ -1,78 +0,0 @@
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestDataParallelism(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_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=["--dp", 2],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
|
||||
def test_update_weight(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/update_weights_from_disk",
|
||||
json={"model_path": DEFAULT_MODEL_NAME_FOR_TEST},
|
||||
)
|
||||
|
||||
# check if the response is 200
|
||||
assert response.status_code == 200
|
||||
|
||||
# pause a few seconds then send again
|
||||
time.sleep(1)
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/update_weights_from_disk",
|
||||
json={"model_path": DEFAULT_MODEL_NAME_FOR_TEST},
|
||||
)
|
||||
|
||||
# check if the response is 200
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_get_memory_pool_size(self):
|
||||
# use `get_server_info` instead since `get_memory_pool_size` is merged into `get_server_info`
|
||||
response = requests.get(self.base_url + "/get_server_info")
|
||||
assert response.status_code == 200
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
response = requests.get(self.base_url + "/get_server_info")
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,438 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import openai
|
||||
import requests
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
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_DRAFT_MODEL_EAGLE,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TARGET_MODEL_EAGLE,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
popen_launch_pd_server,
|
||||
)
|
||||
|
||||
|
||||
class TestDisaggregationAccuracy(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
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 start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"1",
|
||||
]
|
||||
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",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
]
|
||||
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)
|
||||
|
||||
def test_logprob(self):
|
||||
prompt = "The capital of france is "
|
||||
response = requests.post(
|
||||
self.lb_url + "/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {"temperature": 0},
|
||||
"return_logprob": True,
|
||||
"return_input_logprob": True,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
|
||||
j = response.json()
|
||||
completion_tokens = j["meta_info"]["completion_tokens"]
|
||||
input_logprobs = j["meta_info"]["input_token_logprobs"]
|
||||
output_logprobs = j["meta_info"]["output_token_logprobs"]
|
||||
|
||||
assert (
|
||||
len(output_logprobs) == completion_tokens
|
||||
), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
|
||||
assert (
|
||||
len(input_logprobs) > 0
|
||||
), f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
|
||||
|
||||
def test_structured_output(self):
|
||||
json_schema = json.dumps(
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "pattern": "^[\\w]+$"},
|
||||
"population": {"type": "integer"},
|
||||
},
|
||||
"required": ["name", "population"],
|
||||
}
|
||||
)
|
||||
|
||||
# JSON
|
||||
response = requests.post(
|
||||
f"{self.lb_url}/generate",
|
||||
json={
|
||||
"text": "Here is the information of the capital of France in the JSON format.\n",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 64,
|
||||
"json_schema": json_schema,
|
||||
},
|
||||
},
|
||||
)
|
||||
output = response.json()["text"]
|
||||
# ensure the output is a valid JSON
|
||||
json.loads(output)
|
||||
|
||||
def test_first_token_finish(self):
|
||||
client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1")
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.model)
|
||||
eos_token = tokenizer.eos_token_id
|
||||
prompt = "The best programming language for AI is"
|
||||
|
||||
# First token EOS
|
||||
res = client.completions.create(
|
||||
model="dummy", prompt=prompt, logit_bias={eos_token: 42}
|
||||
).model_dump()
|
||||
print(f"{res=}")
|
||||
|
||||
assert res["usage"]["completion_tokens"] == 1, (
|
||||
"Expected completion_tokens to be 1 when first token is EOS, "
|
||||
f"but got {res['usage']['completion_tokens']}"
|
||||
)
|
||||
|
||||
# First token EOS with ignore_eos
|
||||
res = client.completions.create(
|
||||
model="dummy",
|
||||
prompt=prompt,
|
||||
logit_bias={eos_token: 42},
|
||||
extra_body={"ignore_eos": True},
|
||||
).model_dump()
|
||||
print(f"{res=}")
|
||||
|
||||
assert res["usage"]["completion_tokens"] > 1, (
|
||||
"Expected completion_tokens to be greater than 1 when ignore_eos is True, "
|
||||
f"but got {res['usage']['completion_tokens']}"
|
||||
)
|
||||
|
||||
# First token with specified stop token
|
||||
stop_token_id = tokenizer.encode(" hello", add_special_tokens=False)[0]
|
||||
res = client.completions.create(
|
||||
model="dummy",
|
||||
prompt=prompt,
|
||||
logit_bias={stop_token_id: 42},
|
||||
stop=[" hello"],
|
||||
).model_dump()
|
||||
print(f"{res=}")
|
||||
|
||||
assert res["usage"]["completion_tokens"] == 1, (
|
||||
"Expected completion_tokens to be 1 when first token is stop token, "
|
||||
f"but got {res['usage']['completion_tokens']}"
|
||||
)
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
# set DISAGGREGATION_TEST_FAILURE_PROB to simulate failure
|
||||
os.environ["DISAGGREGATION_TEST_FAILURE_PROB"] = "0.05"
|
||||
|
||||
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("DISAGGREGATION_TEST_FAILURE_PROB")
|
||||
super().tearDownClass()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"1",
|
||||
]
|
||||
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",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
]
|
||||
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),
|
||||
)
|
||||
|
||||
# Expect lots of failure but the server cannot crash
|
||||
try:
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Evaluation metrics: {metrics}")
|
||||
except Exception as e:
|
||||
print(f"Test encountered expected errors: {e}")
|
||||
# Check if servers are still healthy
|
||||
try:
|
||||
response = requests.get(self.prefill_url + "/health_generate")
|
||||
assert response.status_code == 200
|
||||
response = requests.get(self.decode_url + "/health_generate")
|
||||
assert response.status_code == 200
|
||||
except Exception as health_check_error:
|
||||
# If health check fails, re-raise the original exception
|
||||
raise e from health_check_error
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeSpec(PDDisaggregationServerBase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = DEFAULT_TARGET_MODEL_EAGLE
|
||||
cls.draft_model = DEFAULT_DRAFT_MODEL_EAGLE
|
||||
cls.spec_args = [
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
cls.draft_model,
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"4",
|
||||
"--speculative-num-draft-tokens",
|
||||
"16",
|
||||
"--cuda-graph-max-bs",
|
||||
"8",
|
||||
]
|
||||
print(f"{cls.base_host=} {cls.lb_port=} {cls.prefill_port=} {cls.decode_port=}")
|
||||
|
||||
# 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",
|
||||
"1",
|
||||
] + cls.spec_args
|
||||
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",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
] + cls.spec_args
|
||||
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=2,
|
||||
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.20)
|
||||
|
||||
|
||||
class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
os.environ["SGLANG_TEST_RETRACT"] = "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_TEST_RETRACT")
|
||||
super().tearDownClass()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp",
|
||||
"1",
|
||||
]
|
||||
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",
|
||||
"1",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
]
|
||||
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()
|
||||
@@ -1,165 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test
|
||||
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_MLA,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_amd_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestDPAttentionDP2TP2(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",
|
||||
"--enable-dp-attention",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-torch-compile",
|
||||
"--torch-compile-max-bs",
|
||||
"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 TestDPRetract(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",
|
||||
"--enable-dp-attention",
|
||||
"--dp",
|
||||
"2",
|
||||
"--max-total-tokens",
|
||||
"4500",
|
||||
"--max-running-requests",
|
||||
"128",
|
||||
"--chunked-prefill-size",
|
||||
"256",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_radix_attention(self):
|
||||
with envs.SGLANG_TEST_RETRACT.override(True):
|
||||
run_radix_attention_test(self.base_url)
|
||||
self.assertIsNone(self.process.poll())
|
||||
|
||||
|
||||
class TestDPAttentionDP2TP2DeepseekV3MTP(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--disable-radix",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-num-steps",
|
||||
"2",
|
||||
"--speculative-eagle-topk",
|
||||
"4",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--dp-size",
|
||||
"2",
|
||||
]
|
||||
if not is_in_amd_ci():
|
||||
other_args += ["--mem-frac", "0.7"]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
|
||||
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(metrics)
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.60)
|
||||
|
||||
server_info = requests.get(self.base_url + "/get_server_info")
|
||||
avg_spec_accept_length = server_info.json()["internal_states"][0][
|
||||
"avg_spec_accept_length"
|
||||
]
|
||||
print(
|
||||
f"###test_gsm8k (deepseek-v3 mtp + dp):\n"
|
||||
f"accuracy={metrics['accuracy']=:.3f}\n"
|
||||
f"{avg_spec_accept_length=:.3f}\n"
|
||||
)
|
||||
self.assertGreater(avg_spec_accept_length, 2.5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,431 +0,0 @@
|
||||
"""Test loading weights from remote instance.
|
||||
|
||||
This test suite simulates loading weights from a remote instance.
|
||||
Rank 0 represents the seed instance, while ranks 1 represents the
|
||||
new instance that needs to loading weights from the seed instance.
|
||||
|
||||
Seed instance must be started in `Server` mode, while the dst instance
|
||||
can be either `Engine` mode or `Server` mode.
|
||||
|
||||
Seed instance does not support concurrently serving multiple dst instances.
|
||||
User has to guarantee that there is only one dst instance trying to load
|
||||
weights from the seed instance at any time.
|
||||
|
||||
"""
|
||||
|
||||
import gc
|
||||
import os
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
from sglang.utils import terminate_process
|
||||
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
|
||||
def verify_params_close(params1, params2, error_msg):
|
||||
"""Verify if two parameter arrays are close enough."""
|
||||
try:
|
||||
assert np.allclose(np.array(params1), np.array(params2)), error_msg
|
||||
except Exception as e:
|
||||
print(f"Parameters not close for {error_msg}")
|
||||
print("Params1:", np.array(params1))
|
||||
print("Params2:", np.array(params2))
|
||||
raise e
|
||||
|
||||
|
||||
def init_process(
|
||||
rank,
|
||||
param_queue,
|
||||
truncate_size,
|
||||
tp_size,
|
||||
model_name,
|
||||
backends,
|
||||
checking_parameters,
|
||||
seed_instance_ip,
|
||||
seed_instance_service_port,
|
||||
seed_instance_group_base_port,
|
||||
event_seed_ready,
|
||||
event_dst_ready_list,
|
||||
remote_instance_loader_backend,
|
||||
):
|
||||
torch.cuda.set_device(rank)
|
||||
|
||||
if rank == 0:
|
||||
init_process_seed(
|
||||
rank,
|
||||
param_queue,
|
||||
truncate_size,
|
||||
model_name,
|
||||
checking_parameters,
|
||||
tp_size,
|
||||
event_seed_ready,
|
||||
event_dst_ready_list,
|
||||
)
|
||||
elif rank in [1, 2]:
|
||||
init_process_dst(
|
||||
rank,
|
||||
param_queue,
|
||||
truncate_size,
|
||||
model_name,
|
||||
seed_instance_ip,
|
||||
seed_instance_service_port,
|
||||
seed_instance_group_base_port,
|
||||
checking_parameters,
|
||||
backends[rank - 1],
|
||||
tp_size,
|
||||
event_seed_ready,
|
||||
event_dst_ready_list,
|
||||
remote_instance_loader_backend,
|
||||
)
|
||||
|
||||
|
||||
def init_process_seed(
|
||||
rank,
|
||||
param_queue,
|
||||
truncate_size,
|
||||
model_name,
|
||||
checking_parameters,
|
||||
tp_size,
|
||||
event_seed_ready,
|
||||
event_dst_ready_list,
|
||||
):
|
||||
# These two environment variables are very important
|
||||
# to avoid unexpected behaviors of CUDA and NCCL.
|
||||
os.environ["NCCL_CUMEM_ENABLE"] = "0"
|
||||
os.environ["NCCL_NVLS_ENABLE"] = "0"
|
||||
|
||||
# Load model and get parameters
|
||||
torch.cuda.set_device(rank)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
url = DEFAULT_URL_FOR_TEST
|
||||
process = popen_launch_server(
|
||||
model_name,
|
||||
url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=(
|
||||
"--base-gpu-id",
|
||||
str(rank),
|
||||
"--tp-size",
|
||||
str(tp_size),
|
||||
"--remote-instance-weight-loader-start-seed-via-transfer-engine",
|
||||
),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
seed_params = []
|
||||
# Get the weights of seed instance for correctness check.
|
||||
for parameter_name in checking_parameters:
|
||||
seed_params.append(
|
||||
requests.get(
|
||||
f"{url}/get_weights_by_name",
|
||||
json={
|
||||
"name": parameter_name,
|
||||
"truncate_size": truncate_size,
|
||||
},
|
||||
).json()
|
||||
)
|
||||
param_queue.put((f"seed_params", seed_params))
|
||||
|
||||
event_seed_ready.set()
|
||||
for i in range(len(event_dst_ready_list)):
|
||||
event_dst_ready_list[i].wait()
|
||||
terminate_process(process)
|
||||
|
||||
|
||||
def init_process_dst(
|
||||
rank,
|
||||
param_queue,
|
||||
truncate_size,
|
||||
model_name,
|
||||
seed_instance_ip,
|
||||
seed_instance_service_port,
|
||||
seed_instance_group_base_port,
|
||||
checking_parameters,
|
||||
backend,
|
||||
tp_size,
|
||||
event_seed_ready,
|
||||
event_dst_ready_list,
|
||||
remote_instance_loader_backend,
|
||||
):
|
||||
torch.cuda.set_device(rank * tp_size)
|
||||
torch.cuda.synchronize()
|
||||
base_gpu_id = rank * tp_size
|
||||
|
||||
event_seed_ready.wait()
|
||||
print(f"rank {rank}, seed ready")
|
||||
for i in range(rank - 1):
|
||||
print(f"rank {rank}, wait dst {i}")
|
||||
event_dst_ready_list[i].wait()
|
||||
|
||||
ports = []
|
||||
for i in range(tp_size):
|
||||
ports.append(seed_instance_group_base_port + (rank - 1) * tp_size + i)
|
||||
|
||||
if backend == "Engine":
|
||||
print(f"[sgl] rank {rank} init engine")
|
||||
engine = sgl.Engine(
|
||||
model_path=model_name,
|
||||
base_gpu_id=base_gpu_id,
|
||||
tp_size=tp_size,
|
||||
cuda_graph_max_bs=2,
|
||||
tokenizer_path=model_name,
|
||||
remote_instance_weight_loader_seed_instance_ip=seed_instance_ip,
|
||||
remote_instance_weight_loader_seed_instance_service_port=seed_instance_service_port,
|
||||
remote_instance_weight_loader_send_weights_group_ports=ports,
|
||||
load_format="remote_instance",
|
||||
remote_instance_weight_loader_backend=remote_instance_loader_backend,
|
||||
remote_instance_weight_loader_start_seed_via_transfer_engine=(
|
||||
remote_instance_loader_backend == "transfer_engine"
|
||||
),
|
||||
)
|
||||
else:
|
||||
host, _, port = DEFAULT_URL_FOR_TEST.rpartition(":")
|
||||
url = ":".join([host, str(int(port) + 10000 + rank)])
|
||||
|
||||
print(f"[sgl] rank {rank} init server on url: {url}")
|
||||
process = popen_launch_server(
|
||||
model_name,
|
||||
url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=(
|
||||
"--base-gpu-id",
|
||||
str(base_gpu_id),
|
||||
"--tp-size",
|
||||
str(tp_size),
|
||||
"--cuda-graph-max-bs",
|
||||
2,
|
||||
"--tokenizer-path",
|
||||
model_name,
|
||||
"--remote-instance-weight-loader-seed-instance-ip",
|
||||
seed_instance_ip,
|
||||
"--remote-instance-weight-loader-seed-instance-service-port",
|
||||
seed_instance_service_port,
|
||||
"--remote-instance-weight-loader-send-weights-group-ports",
|
||||
f"[{','.join(str(port) for port in ports)}]",
|
||||
"--load-format",
|
||||
"remote_instance",
|
||||
"--remote-instance-weight-loader-backend",
|
||||
remote_instance_loader_backend,
|
||||
"--remote-instance-weight-loader-start-seed-via-transfer-engine",
|
||||
),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
event_dst_ready_list[rank - 1].set()
|
||||
|
||||
# Get weights of destination instance loaded from remote instance.
|
||||
dst_params = []
|
||||
for parameter_name in checking_parameters:
|
||||
dst_params.append(
|
||||
engine.get_weights_by_name(parameter_name, truncate_size)
|
||||
if backend == "Engine"
|
||||
else requests.get(
|
||||
f"{url}/get_weights_by_name",
|
||||
json={"name": parameter_name, "truncate_size": truncate_size},
|
||||
).json()
|
||||
)
|
||||
|
||||
param_queue.put((f"sgl_dp_{rank}_dst_params", dst_params))
|
||||
|
||||
# Shutdown the engine or terminate the server process.
|
||||
if backend == "Engine":
|
||||
engine.shutdown()
|
||||
else:
|
||||
terminate_process(process)
|
||||
|
||||
|
||||
def test_load_weights_from_remote_instance(
|
||||
tp_size,
|
||||
dp_size,
|
||||
model_name,
|
||||
backends,
|
||||
truncate_size,
|
||||
checking_parameters,
|
||||
seed_instance_ip,
|
||||
seed_instance_service_port,
|
||||
seed_instance_group_base_port,
|
||||
remote_instance_loader_backend,
|
||||
):
|
||||
print(
|
||||
f"Testing model: {model_name} tp_size: {tp_size}, dp_size: {dp_size} backend: {backends} remote_instance_loader_backend: {remote_instance_loader_backend}"
|
||||
)
|
||||
param_queue = mp.Queue()
|
||||
results = {}
|
||||
event_seed_ready = mp.Event()
|
||||
event_dst_ready_list = []
|
||||
for i in range(dp_size):
|
||||
event_dst_ready = mp.Event()
|
||||
event_dst_ready_list.append(event_dst_ready)
|
||||
|
||||
context = mp.spawn(
|
||||
init_process,
|
||||
args=(
|
||||
param_queue,
|
||||
truncate_size,
|
||||
tp_size,
|
||||
model_name,
|
||||
backends,
|
||||
checking_parameters,
|
||||
seed_instance_ip,
|
||||
seed_instance_service_port,
|
||||
seed_instance_group_base_port,
|
||||
event_seed_ready,
|
||||
event_dst_ready_list,
|
||||
remote_instance_loader_backend,
|
||||
),
|
||||
nprocs=1 + dp_size,
|
||||
join=False,
|
||||
)
|
||||
|
||||
while len(results) < (1 + dp_size):
|
||||
try:
|
||||
key, value = param_queue.get(timeout=5)
|
||||
results[key] = value
|
||||
except Exception as e:
|
||||
if all(not p.is_alive() for p in context.processes):
|
||||
break
|
||||
|
||||
context.join()
|
||||
|
||||
if len(results) != (1 + dp_size):
|
||||
raise RuntimeError(
|
||||
f"Expected {(1 + dp_size)} parameters but got {len(results)}"
|
||||
)
|
||||
|
||||
params = {
|
||||
"seed": results.get("seed_params"),
|
||||
"sgl_dp_1_dest": results.get("sgl_dp_1_dst_params"),
|
||||
}
|
||||
|
||||
if dp_size == 2:
|
||||
dp2_params = {
|
||||
"sgl_dp_2_dest": results.get("sgl_dp_2_dst_params"),
|
||||
}
|
||||
assert all(v is not None for v in dp2_params.values())
|
||||
params.update(dp2_params)
|
||||
|
||||
# Check the correctness of weights loaded from remote instance
|
||||
# by verifying the weights of seed instance and destination instance.
|
||||
for i in range(len(params["seed"])):
|
||||
verify_params_close(
|
||||
params["seed"][i],
|
||||
params["sgl_dp_1_dest"][i],
|
||||
f"sgl_dp_1_dst_params rank {i}",
|
||||
)
|
||||
|
||||
if dp_size == 2:
|
||||
verify_params_close(
|
||||
params["seed"][i],
|
||||
params["sgl_dp_2_dest"][i],
|
||||
f"sgl_dp_2_dst_params rank {i}",
|
||||
)
|
||||
|
||||
# Delete the context and close the parameter queue.
|
||||
del context
|
||||
param_queue.close()
|
||||
param_queue.join_thread()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
class TestLoadWeightsFromRemoteInstance(CustomTestCase):
|
||||
|
||||
def test_load_weights_from_remote_instance(self):
|
||||
|
||||
assert torch.cuda.device_count() >= 2, "At least 2 GPUs are required"
|
||||
# test_suits : tp, dp, model_name, backend, dst_instance_id
|
||||
if is_in_ci():
|
||||
mode = random.choice(["Engine", "Server"])
|
||||
remote_instance_loader_backend = random.choice(["nccl", "transfer_engine"])
|
||||
test_suits = [
|
||||
(
|
||||
1,
|
||||
1,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
[mode],
|
||||
remote_instance_loader_backend,
|
||||
),
|
||||
]
|
||||
else:
|
||||
test_suits = [
|
||||
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Engine"], "nccl"),
|
||||
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Server"], "nccl"),
|
||||
(2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Engine", "Server"], "nccl"),
|
||||
(
|
||||
1,
|
||||
1,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
["Engine"],
|
||||
"transfer_engine",
|
||||
),
|
||||
(
|
||||
1,
|
||||
1,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
["Server"],
|
||||
"transfer_engine",
|
||||
),
|
||||
(
|
||||
2,
|
||||
2,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
["Engine", "Server"],
|
||||
"transfer_engine",
|
||||
),
|
||||
]
|
||||
|
||||
truncate_size = 10
|
||||
checking_parameters = [
|
||||
"model.embed_tokens.weight",
|
||||
"model.layers.0.input_layernorm.weight",
|
||||
"model.layers.1.self_attn.q_proj.weight",
|
||||
"model.layers.2.self_attn.k_proj.weight",
|
||||
"model.layers.3.self_attn.v_proj.weight",
|
||||
"model.layers.4.self_attn.o_proj.weight",
|
||||
"model.layers.5.mlp.gate_proj.weight",
|
||||
"model.layers.6.mlp.up_proj.weight",
|
||||
"model.layers.7.mlp.down_proj.weight",
|
||||
"model.layers.8.post_attention_layernorm.weight",
|
||||
"model.norm.weight",
|
||||
]
|
||||
|
||||
for (
|
||||
tp_size,
|
||||
dp_size,
|
||||
model_name,
|
||||
backends,
|
||||
remote_instance_loader_backend,
|
||||
) in test_suits:
|
||||
test_load_weights_from_remote_instance(
|
||||
tp_size,
|
||||
dp_size,
|
||||
model_name,
|
||||
backends,
|
||||
truncate_size,
|
||||
checking_parameters,
|
||||
"127.0.0.1",
|
||||
DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000,
|
||||
60000,
|
||||
remote_instance_loader_backend,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user