Fix external_models import path and migrate model loading tests (#16458)
This commit is contained in:
@@ -0,0 +1,32 @@
|
||||
import unittest
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=30, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=45, suite="stage-b-test-small-1-gpu")
|
||||
|
||||
|
||||
class TestExternalModels(CustomTestCase):
|
||||
def test_external_model(self):
|
||||
envs.SGLANG_EXTERNAL_MODEL_PACKAGE.set("sglang.test.external_models")
|
||||
envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.set("sglang.test.external_models")
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = "Qwen/Qwen2-VL-2B-Instruct"
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path,
|
||||
cuda_graph_max_bs=1,
|
||||
max_total_tokens=64,
|
||||
enable_multimodal=True,
|
||||
)
|
||||
out = engine.generate(prompt)["text"]
|
||||
engine.shutdown()
|
||||
|
||||
self.assertGreater(len(out), 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,355 @@
|
||||
"""
|
||||
Unit tests for ModelOpt export functionality in SGLang.
|
||||
|
||||
These tests verify the integration of ModelOpt export API with SGLang's model loading
|
||||
and quantization workflow.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.device_config import DeviceConfig
|
||||
from sglang.srt.configs.load_config import LoadConfig
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.model_loader.loader import ModelOptModelLoader
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
|
||||
|
||||
# Note: PYTHONPATH=python should be set when running tests
|
||||
|
||||
# Check if modelopt is available
|
||||
try:
|
||||
import modelopt # noqa: F401
|
||||
|
||||
MODELOPT_AVAILABLE = True
|
||||
except ImportError:
|
||||
MODELOPT_AVAILABLE = False
|
||||
|
||||
|
||||
class TestModelOptExport(unittest.TestCase):
|
||||
"""Test suite for ModelOpt export functionality."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
# Mock distributed functionality to avoid initialization errors
|
||||
self.mock_tp_rank = patch(
|
||||
"sglang.srt.distributed.parallel_state.get_tensor_model_parallel_rank",
|
||||
return_value=0,
|
||||
)
|
||||
self.mock_tp_rank.start()
|
||||
|
||||
self.mock_rank0_log = patch("sglang.srt.model_loader.loader.rank0_log")
|
||||
self.mock_rank0_log.start()
|
||||
|
||||
# Mock logger to avoid issues
|
||||
self.mock_logger = patch("sglang.srt.model_loader.loader.logger")
|
||||
self.mock_logger.start()
|
||||
|
||||
# Mock all distributed functions that might be called
|
||||
self.mock_get_tp_group = patch(
|
||||
"sglang.srt.distributed.parallel_state.get_tp_group"
|
||||
)
|
||||
self.mock_get_tp_group.start()
|
||||
|
||||
# Mock model parallel initialization check
|
||||
self.mock_mp_is_initialized = patch(
|
||||
"sglang.srt.distributed.parallel_state.model_parallel_is_initialized",
|
||||
return_value=True,
|
||||
)
|
||||
self.mock_mp_is_initialized.start()
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.export_dir = os.path.join(self.temp_dir, "exported_model")
|
||||
self.checkpoint_dir = os.path.join(self.temp_dir, "checkpoint")
|
||||
|
||||
# Mock model
|
||||
self.mock_model = Mock(spec=torch.nn.Module)
|
||||
self.mock_model.device = torch.device("cuda:0")
|
||||
|
||||
# Mock tokenizer
|
||||
self.mock_tokenizer = Mock()
|
||||
|
||||
# Mock quantization config
|
||||
self.mock_quant_cfg = Mock()
|
||||
|
||||
# Create ModelOptModelLoader instance
|
||||
self.load_config = LoadConfig()
|
||||
self.model_loader = ModelOptModelLoader(self.load_config)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up test fixtures."""
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
# Stop mocks
|
||||
self.mock_tp_rank.stop()
|
||||
self.mock_rank0_log.stop()
|
||||
self.mock_logger.stop()
|
||||
self.mock_get_tp_group.stop()
|
||||
self.mock_mp_is_initialized.stop()
|
||||
|
||||
def _create_mock_export_files(self, export_dir: str):
|
||||
"""Create mock export files for testing validation."""
|
||||
os.makedirs(export_dir, exist_ok=True)
|
||||
|
||||
# Create config.json
|
||||
config = {
|
||||
"model_type": "test_model",
|
||||
"architectures": ["TestModel"],
|
||||
"quantization_config": {
|
||||
"quant_method": "modelopt",
|
||||
"bits": 8,
|
||||
},
|
||||
}
|
||||
with open(os.path.join(export_dir, "config.json"), "w") as f:
|
||||
json.dump(config, f)
|
||||
|
||||
# Create tokenizer_config.json
|
||||
tokenizer_config = {"tokenizer_class": "TestTokenizer"}
|
||||
with open(os.path.join(export_dir, "tokenizer_config.json"), "w") as f:
|
||||
json.dump(tokenizer_config, f)
|
||||
|
||||
# Create model file
|
||||
with open(os.path.join(export_dir, "model.safetensors"), "w") as f:
|
||||
f.write("mock_model_data")
|
||||
|
||||
@unittest.skipIf(not MODELOPT_AVAILABLE, "nvidia-modelopt not available")
|
||||
@patch("sglang.srt.model_loader.loader.os.makedirs")
|
||||
@patch("modelopt.torch.export.export_hf_checkpoint")
|
||||
def test_export_modelopt_checkpoint_success(self, mock_export, mock_makedirs):
|
||||
"""Test successful model export."""
|
||||
# Arrange
|
||||
mock_export.return_value = None
|
||||
mock_makedirs.return_value = None
|
||||
|
||||
# Act
|
||||
self.model_loader._export_modelopt_checkpoint(self.mock_model, self.export_dir)
|
||||
|
||||
# Assert
|
||||
mock_makedirs.assert_called_once_with(self.export_dir, exist_ok=True)
|
||||
mock_export.assert_called_once_with(self.mock_model, export_dir=self.export_dir)
|
||||
|
||||
@unittest.skipIf(not MODELOPT_AVAILABLE, "nvidia-modelopt not available")
|
||||
@patch("modelopt.torch.opt.restore")
|
||||
@patch("modelopt.torch.quantization.utils.is_quantized")
|
||||
def test_setup_quantization_with_export_from_checkpoint(
|
||||
self, mock_is_quantized, mock_restore
|
||||
):
|
||||
"""Test export functionality when restoring from checkpoint."""
|
||||
# Arrange
|
||||
mock_is_quantized.return_value = False
|
||||
mock_restore.return_value = None
|
||||
|
||||
with patch.object(
|
||||
self.model_loader, "_export_modelopt_checkpoint"
|
||||
) as mock_export:
|
||||
# Act
|
||||
self.model_loader._setup_modelopt_quantization(
|
||||
self.mock_model,
|
||||
self.mock_tokenizer,
|
||||
self.mock_quant_cfg,
|
||||
quantized_ckpt_restore_path=self.checkpoint_dir,
|
||||
export_path=self.export_dir,
|
||||
)
|
||||
|
||||
# Assert
|
||||
mock_restore.assert_called_once_with(self.mock_model, self.checkpoint_dir)
|
||||
mock_export.assert_called_once_with(self.mock_model, self.export_dir, None)
|
||||
|
||||
@unittest.skipIf(not MODELOPT_AVAILABLE, "nvidia-modelopt not available")
|
||||
@patch("modelopt.torch.quantization.quantize")
|
||||
@patch("modelopt.torch.quantization.print_quant_summary")
|
||||
@patch("modelopt.torch.quantization.utils.is_quantized")
|
||||
@patch("modelopt.torch.utils.dataset_utils.get_dataset_dataloader")
|
||||
@patch("modelopt.torch.utils.dataset_utils.create_forward_loop")
|
||||
def test_setup_quantization_with_export_after_calibration(
|
||||
self,
|
||||
mock_create_loop,
|
||||
mock_get_dataloader,
|
||||
mock_is_quantized,
|
||||
mock_print_summary,
|
||||
mock_quantize,
|
||||
):
|
||||
"""Test export functionality after calibration-based quantization."""
|
||||
# Arrange
|
||||
mock_is_quantized.return_value = False
|
||||
mock_dataloader = Mock()
|
||||
mock_get_dataloader.return_value = mock_dataloader
|
||||
mock_calibrate_loop = Mock()
|
||||
mock_create_loop.return_value = mock_calibrate_loop
|
||||
mock_quantize.return_value = None
|
||||
mock_print_summary.return_value = None
|
||||
|
||||
with patch.object(
|
||||
self.model_loader, "_export_modelopt_checkpoint"
|
||||
) as mock_export:
|
||||
# Act
|
||||
self.model_loader._setup_modelopt_quantization(
|
||||
self.mock_model,
|
||||
self.mock_tokenizer,
|
||||
self.mock_quant_cfg,
|
||||
export_path=self.export_dir,
|
||||
)
|
||||
|
||||
# Assert
|
||||
mock_quantize.assert_called_once_with(
|
||||
self.mock_model, self.mock_quant_cfg, forward_loop=mock_calibrate_loop
|
||||
)
|
||||
mock_export.assert_called_once_with(self.mock_model, self.export_dir, None)
|
||||
|
||||
@unittest.skipIf(not MODELOPT_AVAILABLE, "nvidia-modelopt not available")
|
||||
def test_setup_quantization_without_export(self):
|
||||
"""Test quantization setup without export path specified."""
|
||||
with patch("modelopt.torch.quantization.utils.is_quantized", return_value=True):
|
||||
# Act
|
||||
with patch.object(
|
||||
self.model_loader, "_export_modelopt_checkpoint"
|
||||
) as mock_export:
|
||||
self.model_loader._setup_modelopt_quantization(
|
||||
self.mock_model,
|
||||
self.mock_tokenizer,
|
||||
self.mock_quant_cfg,
|
||||
export_path=None, # No export path
|
||||
)
|
||||
|
||||
# Assert
|
||||
mock_export.assert_not_called()
|
||||
|
||||
def test_quantize_and_serve_config_validation(self):
|
||||
"""Test that quantize_and_serve is properly disabled."""
|
||||
# Test that quantize-and-serve mode raises NotImplementedError
|
||||
with self.assertRaises(NotImplementedError) as context:
|
||||
ModelConfig(
|
||||
model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
quantization="modelopt_fp8",
|
||||
quantize_and_serve=True,
|
||||
)
|
||||
|
||||
# Verify the error message contains helpful instructions
|
||||
error_msg = str(context.exception)
|
||||
self.assertIn("disabled due to compatibility issues", error_msg)
|
||||
self.assertIn("separate quantize-then-deploy workflow", error_msg)
|
||||
|
||||
# Test invalid configuration - no quantization
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ModelConfig(
|
||||
model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
quantize_and_serve=True,
|
||||
)
|
||||
self.assertIn("requires ModelOpt quantization", str(context.exception))
|
||||
|
||||
@unittest.skipIf(not MODELOPT_AVAILABLE, "nvidia-modelopt not available")
|
||||
def test_standard_workflow_selection(self):
|
||||
"""Test that standard workflow is selected by default."""
|
||||
with patch(
|
||||
"modelopt.torch.quantization.utils.is_quantized", return_value=False
|
||||
):
|
||||
with patch.object(
|
||||
self.model_loader, "_standard_quantization_workflow"
|
||||
) as mock_standard:
|
||||
with patch.object(self.model_loader, "_load_modelopt_base_model"):
|
||||
mock_standard.return_value = Mock()
|
||||
|
||||
# Create model config without quantize_and_serve
|
||||
model_config = ModelConfig(
|
||||
model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
quantization="modelopt_fp8",
|
||||
quantize_and_serve=False,
|
||||
)
|
||||
device_config = DeviceConfig()
|
||||
|
||||
# Act
|
||||
self.model_loader.load_model(
|
||||
model_config=model_config,
|
||||
device_config=device_config,
|
||||
)
|
||||
|
||||
# Assert
|
||||
mock_standard.assert_called_once_with(model_config, device_config)
|
||||
|
||||
def _get_export_info(self, export_dir: str) -> dict:
|
||||
"""Get information about an exported model."""
|
||||
if not self._validate_export(export_dir):
|
||||
return None
|
||||
|
||||
try:
|
||||
config_path = os.path.join(export_dir, "config.json")
|
||||
with open(config_path, "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
return {
|
||||
"model_type": config.get("model_type", "unknown"),
|
||||
"architectures": config.get("architectures", []),
|
||||
"quantization_config": config.get("quantization_config", {}),
|
||||
"export_dir": export_dir,
|
||||
}
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@unittest.skipIf(not MODELOPT_AVAILABLE, "nvidia-modelopt not available")
|
||||
class TestModelOptExportIntegration(unittest.TestCase):
|
||||
"""Integration tests for ModelOpt export with full model loading workflow."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up integration test fixtures."""
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.export_dir = os.path.join(self.temp_dir, "exported_model")
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up integration test fixtures."""
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.get_model_architecture")
|
||||
@patch("transformers.AutoTokenizer.from_pretrained")
|
||||
@patch("transformers.AutoModelForCausalLM.from_pretrained")
|
||||
def test_full_workflow_with_export(self, mock_model, mock_tokenizer, mock_arch):
|
||||
"""Test the complete workflow from model config to export."""
|
||||
# Arrange
|
||||
mock_arch.return_value = ("TestModel", "TestConfig")
|
||||
mock_tokenizer.return_value = Mock()
|
||||
mock_model.return_value = Mock(spec=torch.nn.Module)
|
||||
|
||||
model_config = ModelConfig(
|
||||
model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
modelopt_quant="fp8",
|
||||
modelopt_export_path=self.export_dir,
|
||||
)
|
||||
|
||||
load_config = LoadConfig()
|
||||
device_config = DeviceConfig()
|
||||
|
||||
# Mock the quantization and export process
|
||||
with patch.object(
|
||||
ModelOptModelLoader, "_setup_modelopt_quantization"
|
||||
) as mock_setup:
|
||||
with patch.object(
|
||||
ModelOptModelLoader, "_load_modelopt_base_model"
|
||||
) as mock_load_base:
|
||||
mock_load_base.return_value = mock_model.return_value
|
||||
|
||||
# Act
|
||||
model_loader = ModelOptModelLoader(load_config)
|
||||
result = model_loader.load_model(
|
||||
model_config=model_config,
|
||||
device_config=device_config,
|
||||
)
|
||||
|
||||
# Assert
|
||||
self.assertIsNotNone(result)
|
||||
mock_setup.assert_called_once()
|
||||
# Verify export_path was passed to setup
|
||||
args, kwargs = mock_setup.call_args
|
||||
self.assertEqual(kwargs.get("export_path"), self.export_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,564 @@
|
||||
"""
|
||||
Unit tests for ModelOptModelLoader class.
|
||||
|
||||
This test module verifies the functionality of ModelOptModelLoader, which
|
||||
applies NVIDIA Model Optimizer quantization to models during loading.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.srt.configs.device_config import DeviceConfig
|
||||
from sglang.srt.configs.load_config import LoadConfig
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
|
||||
from sglang.srt.model_loader.loader import ModelOptModelLoader
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# Note: PYTHONPATH=python should be set when running tests
|
||||
|
||||
# Constants for calibration parameters to avoid hard-coded values
|
||||
CALIBRATION_BATCH_SIZE = 36
|
||||
CALIBRATION_NUM_SAMPLES = 512
|
||||
DEFAULT_DEVICE = "cuda:0"
|
||||
|
||||
register_cuda_ci(est_time=11, suite="stage-b-test-small-1-gpu")
|
||||
|
||||
|
||||
class TestModelOptModelLoader(CustomTestCase):
|
||||
"""Test cases for ModelOptModelLoader functionality."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
# Mock distributed functionality to avoid initialization errors
|
||||
self.mock_tp_rank = patch(
|
||||
"sglang.srt.distributed.parallel_state.get_tensor_model_parallel_rank",
|
||||
return_value=0,
|
||||
)
|
||||
self.mock_tp_rank.start()
|
||||
|
||||
self.mock_rank0_log = patch("sglang.srt.model_loader.loader.rank0_log")
|
||||
self.mock_rank0_log.start()
|
||||
|
||||
# Mock logger to avoid issues
|
||||
self.mock_logger = patch("sglang.srt.model_loader.loader.logger")
|
||||
self.mock_logger.start()
|
||||
|
||||
# Mock all distributed functions that might be called
|
||||
self.mock_get_tp_group = patch(
|
||||
"sglang.srt.distributed.parallel_state.get_tp_group"
|
||||
)
|
||||
self.mock_get_tp_group.start()
|
||||
|
||||
# Mock model parallel initialization check
|
||||
self.mock_mp_is_initialized = patch(
|
||||
"sglang.srt.distributed.parallel_state.model_parallel_is_initialized",
|
||||
return_value=True,
|
||||
)
|
||||
self.mock_mp_is_initialized.start()
|
||||
|
||||
self.model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
|
||||
self.load_config = LoadConfig()
|
||||
self.device_config = DeviceConfig(device="cuda")
|
||||
|
||||
# Create a basic model config with unified quantization flag
|
||||
self.model_config = ModelConfig(
|
||||
model_path=self.model_path,
|
||||
quantization="modelopt_fp8", # Use unified quantization approach
|
||||
)
|
||||
|
||||
# Also create a unified quantization config for new tests
|
||||
self.unified_model_config = ModelConfig(
|
||||
model_path=self.model_path, quantization="modelopt_fp8"
|
||||
)
|
||||
|
||||
# Mock base model
|
||||
self.mock_base_model = MagicMock(spec=nn.Module)
|
||||
self.mock_base_model.eval.return_value = self.mock_base_model
|
||||
self.mock_base_model.device = (
|
||||
DEFAULT_DEVICE # Add device attribute for calibration tests
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up test fixtures."""
|
||||
# Stop mocks
|
||||
self.mock_tp_rank.stop()
|
||||
self.mock_rank0_log.stop()
|
||||
self.mock_logger.stop()
|
||||
self.mock_get_tp_group.stop()
|
||||
self.mock_mp_is_initialized.stop()
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.QUANT_CFG_CHOICES", QUANT_CFG_CHOICES)
|
||||
@patch("sglang.srt.model_loader.loader.logger")
|
||||
def test_successful_fp8_quantization(self, mock_logger):
|
||||
"""Test successful FP8 quantization workflow."""
|
||||
|
||||
# Create loader instance
|
||||
loader = ModelOptModelLoader(self.load_config)
|
||||
|
||||
# Mock modelopt modules
|
||||
mock_mtq = MagicMock()
|
||||
|
||||
# Configure mtq mock with FP8_DEFAULT_CFG
|
||||
mock_fp8_cfg = MagicMock()
|
||||
mock_mtq.FP8_DEFAULT_CFG = mock_fp8_cfg
|
||||
mock_mtq.quantize.return_value = self.mock_base_model
|
||||
mock_mtq.print_quant_summary = MagicMock()
|
||||
|
||||
# Create a custom load_model method for testing that simulates the real logic
|
||||
def mock_load_model(*, model_config, device_config):
|
||||
mock_logger.info("ModelOptModelLoader: Loading base model...")
|
||||
|
||||
# Simulate loading base model (this is already mocked)
|
||||
model = self.mock_base_model
|
||||
|
||||
# Simulate the quantization config lookup
|
||||
quant_choice_str = model_config._get_modelopt_quant_type()
|
||||
quant_cfg_name = QUANT_CFG_CHOICES.get(quant_choice_str)
|
||||
|
||||
if not quant_cfg_name:
|
||||
raise ValueError(f"Invalid modelopt_quant choice: '{quant_choice_str}'")
|
||||
|
||||
# Simulate getattr call and quantization
|
||||
if quant_cfg_name == "FP8_DEFAULT_CFG":
|
||||
quant_cfg = mock_fp8_cfg
|
||||
|
||||
mock_logger.info(
|
||||
f"Quantizing model with ModelOpt using config attribute: mtq.{quant_cfg_name}"
|
||||
)
|
||||
|
||||
# Simulate mtq.quantize call
|
||||
quantized_model = mock_mtq.quantize(model, quant_cfg, forward_loop=None)
|
||||
mock_logger.info("Model successfully quantized with ModelOpt.")
|
||||
|
||||
# Simulate print_quant_summary call
|
||||
mock_mtq.print_quant_summary(quantized_model)
|
||||
|
||||
return quantized_model.eval()
|
||||
|
||||
return model.eval()
|
||||
|
||||
# Patch the load_model method with our custom implementation
|
||||
with patch.object(loader, "load_model", side_effect=mock_load_model):
|
||||
# Execute the load_model method
|
||||
result_model = loader.load_model(
|
||||
model_config=self.model_config, device_config=self.device_config
|
||||
)
|
||||
|
||||
# Verify the quantization process
|
||||
mock_mtq.quantize.assert_called_once_with(
|
||||
self.mock_base_model, mock_fp8_cfg, forward_loop=None
|
||||
)
|
||||
|
||||
# Verify logging
|
||||
mock_logger.info.assert_any_call(
|
||||
"ModelOptModelLoader: Loading base model..."
|
||||
)
|
||||
mock_logger.info.assert_any_call(
|
||||
"Quantizing model with ModelOpt using config attribute: mtq.FP8_DEFAULT_CFG"
|
||||
)
|
||||
mock_logger.info.assert_any_call(
|
||||
"Model successfully quantized with ModelOpt."
|
||||
)
|
||||
|
||||
# Verify print_quant_summary was called
|
||||
mock_mtq.print_quant_summary.assert_called_once_with(self.mock_base_model)
|
||||
|
||||
# Verify eval() was called on the returned model
|
||||
self.mock_base_model.eval.assert_called()
|
||||
|
||||
# Verify we get back the expected model
|
||||
self.assertEqual(result_model, self.mock_base_model)
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.logger")
|
||||
def test_missing_modelopt_import(self, mock_logger):
|
||||
"""Test error handling when modelopt library is not available."""
|
||||
|
||||
loader = ModelOptModelLoader(self.load_config)
|
||||
|
||||
# Mock the base model loader method
|
||||
with patch.object(
|
||||
loader, "_load_modelopt_base_model", return_value=self.mock_base_model
|
||||
):
|
||||
# Simulate missing modelopt by making import fail
|
||||
original_import = __import__
|
||||
|
||||
def mock_import(name, *args, **kwargs):
|
||||
if name.startswith("modelopt"):
|
||||
raise ImportError("No module named 'modelopt'")
|
||||
# Return default import behavior for other modules
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=mock_import):
|
||||
# Expect ImportError to be raised and logged
|
||||
with self.assertRaises(ImportError):
|
||||
loader.load_model(
|
||||
model_config=self.model_config, device_config=self.device_config
|
||||
)
|
||||
|
||||
# Verify error logging
|
||||
mock_logger.error.assert_called_with(
|
||||
"NVIDIA Model Optimizer (modelopt) library not found. "
|
||||
"Please install it to use ModelOpt quantization."
|
||||
)
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.QUANT_CFG_CHOICES", QUANT_CFG_CHOICES)
|
||||
@patch("sglang.srt.model_loader.loader.AutoTokenizer")
|
||||
@patch("sglang.srt.model_loader.loader.logger")
|
||||
def test_calibration_workflow_integration(self, mock_logger, mock_auto_tokenizer):
|
||||
"""Test end-to-end calibration workflow integration."""
|
||||
|
||||
loader = ModelOptModelLoader(self.load_config)
|
||||
|
||||
# Mock tokenizer
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_tokenizer.padding_side = "right"
|
||||
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
|
||||
|
||||
# Mock modelopt modules
|
||||
mock_mtq = MagicMock()
|
||||
mock_mto = MagicMock()
|
||||
mock_dataset_utils = MagicMock()
|
||||
|
||||
# Configure quantization config
|
||||
mock_fp8_cfg = MagicMock()
|
||||
mock_mtq.FP8_DEFAULT_CFG = mock_fp8_cfg
|
||||
|
||||
# Configure dataset utilities
|
||||
mock_calib_dataloader = MagicMock()
|
||||
mock_calibrate_loop = MagicMock()
|
||||
mock_dataset_utils.get_dataset_dataloader.return_value = mock_calib_dataloader
|
||||
mock_dataset_utils.create_forward_loop.return_value = mock_calibrate_loop
|
||||
|
||||
# Configure model as not quantized initially
|
||||
mock_is_quantized = MagicMock(return_value=False)
|
||||
|
||||
with patch.object(
|
||||
loader, "_load_modelopt_base_model", return_value=self.mock_base_model
|
||||
):
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"modelopt": MagicMock(),
|
||||
"modelopt.torch": MagicMock(),
|
||||
"modelopt.torch.opt": mock_mto,
|
||||
"modelopt.torch.quantization": mock_mtq,
|
||||
"modelopt.torch.quantization.utils": MagicMock(
|
||||
is_quantized=mock_is_quantized
|
||||
),
|
||||
"modelopt.torch.utils": MagicMock(),
|
||||
"modelopt.torch.utils.dataset_utils": mock_dataset_utils,
|
||||
},
|
||||
):
|
||||
# Execute the load_model method to test the full workflow
|
||||
result_model = loader.load_model(
|
||||
model_config=self.model_config, device_config=self.device_config
|
||||
)
|
||||
|
||||
# Verify the model loading was successful
|
||||
self.assertEqual(result_model, self.mock_base_model)
|
||||
|
||||
# Verify key calibration components were used
|
||||
# Note: We can't easily verify the exact calls due to dynamic imports,
|
||||
# but we can verify the workflow completed successfully
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.QUANT_CFG_CHOICES", QUANT_CFG_CHOICES)
|
||||
@patch("sglang.srt.model_loader.loader.AutoTokenizer")
|
||||
@patch("sglang.srt.model_loader.loader.logger")
|
||||
def test_quantized_checkpoint_restore(self, mock_logger, mock_auto_tokenizer):
|
||||
"""Test restoring from a quantized checkpoint."""
|
||||
|
||||
# Create model config with checkpoint restore path
|
||||
config_with_restore = ModelConfig(
|
||||
model_path=self.model_path,
|
||||
quantization="modelopt_fp8",
|
||||
)
|
||||
|
||||
# Create load config with checkpoint restore path
|
||||
load_config_with_restore = LoadConfig(
|
||||
modelopt_checkpoint_restore_path="/path/to/quantized/checkpoint"
|
||||
)
|
||||
|
||||
loader = ModelOptModelLoader(load_config_with_restore)
|
||||
|
||||
# Mock tokenizer
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
|
||||
|
||||
# Mock modelopt modules
|
||||
mock_mtq = MagicMock()
|
||||
mock_mto = MagicMock()
|
||||
|
||||
# Configure quantization config
|
||||
mock_fp8_cfg = MagicMock()
|
||||
mock_mtq.FP8_DEFAULT_CFG = mock_fp8_cfg
|
||||
|
||||
# Configure model as not quantized initially
|
||||
mock_is_quantized = MagicMock(return_value=False)
|
||||
|
||||
with patch.object(
|
||||
loader, "_load_modelopt_base_model", return_value=self.mock_base_model
|
||||
):
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"modelopt": MagicMock(),
|
||||
"modelopt.torch": MagicMock(),
|
||||
"modelopt.torch.opt": mock_mto,
|
||||
"modelopt.torch.quantization": mock_mtq,
|
||||
"modelopt.torch.quantization.utils": MagicMock(
|
||||
is_quantized=mock_is_quantized
|
||||
),
|
||||
},
|
||||
):
|
||||
with patch.object(loader, "_setup_modelopt_quantization") as mock_setup:
|
||||
# Mock the _setup_modelopt_quantization to simulate checkpoint restore
|
||||
def mock_setup_quantization(
|
||||
model,
|
||||
tokenizer,
|
||||
quant_cfg,
|
||||
quantized_ckpt_restore_path=None,
|
||||
**kwargs,
|
||||
):
|
||||
if quantized_ckpt_restore_path:
|
||||
mock_mto.restore(model, quantized_ckpt_restore_path)
|
||||
print(
|
||||
f"Restored quantized model from {quantized_ckpt_restore_path}"
|
||||
)
|
||||
return
|
||||
|
||||
mock_setup.side_effect = mock_setup_quantization
|
||||
|
||||
# Execute the load_model method
|
||||
result_model = loader.load_model(
|
||||
model_config=config_with_restore,
|
||||
device_config=self.device_config,
|
||||
)
|
||||
|
||||
# Verify the setup was called with restore path
|
||||
mock_setup.assert_called_once()
|
||||
call_args = mock_setup.call_args
|
||||
# Check that the restore path was passed correctly
|
||||
self.assertIn("quantized_ckpt_restore_path", call_args[1])
|
||||
self.assertEqual(
|
||||
call_args[1]["quantized_ckpt_restore_path"],
|
||||
"/path/to/quantized/checkpoint",
|
||||
)
|
||||
|
||||
# Verify restore was called
|
||||
mock_mto.restore.assert_called_once_with(
|
||||
self.mock_base_model, "/path/to/quantized/checkpoint"
|
||||
)
|
||||
|
||||
# Verify we get the expected model back
|
||||
self.assertEqual(result_model, self.mock_base_model)
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.QUANT_CFG_CHOICES", QUANT_CFG_CHOICES)
|
||||
@patch("sglang.srt.model_loader.loader.AutoTokenizer")
|
||||
@patch("sglang.srt.model_loader.loader.logger")
|
||||
def test_quantized_checkpoint_save(self, mock_logger, mock_auto_tokenizer):
|
||||
"""Test saving quantized checkpoint after calibration."""
|
||||
|
||||
# Create model config with checkpoint save path
|
||||
config_with_save = ModelConfig(
|
||||
model_path=self.model_path,
|
||||
quantization="modelopt_fp8",
|
||||
)
|
||||
|
||||
# Create load config with checkpoint save path
|
||||
load_config_with_save = LoadConfig(
|
||||
modelopt_checkpoint_save_path="/path/to/save/checkpoint"
|
||||
)
|
||||
|
||||
loader = ModelOptModelLoader(load_config_with_save)
|
||||
|
||||
# Mock tokenizer
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
|
||||
|
||||
# Mock modelopt modules
|
||||
mock_mtq = MagicMock()
|
||||
mock_mto = MagicMock()
|
||||
mock_dataset_utils = MagicMock()
|
||||
|
||||
# Configure quantization config
|
||||
mock_fp8_cfg = MagicMock()
|
||||
mock_mtq.FP8_DEFAULT_CFG = mock_fp8_cfg
|
||||
|
||||
# Configure model as not quantized initially
|
||||
mock_is_quantized = MagicMock(return_value=False)
|
||||
|
||||
with patch.object(
|
||||
loader, "_load_modelopt_base_model", return_value=self.mock_base_model
|
||||
):
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"modelopt": MagicMock(),
|
||||
"modelopt.torch": MagicMock(),
|
||||
"modelopt.torch.opt": mock_mto,
|
||||
"modelopt.torch.quantization": mock_mtq,
|
||||
"modelopt.torch.quantization.utils": MagicMock(
|
||||
is_quantized=mock_is_quantized
|
||||
),
|
||||
"modelopt.torch.utils": MagicMock(),
|
||||
"modelopt.torch.utils.dataset_utils": mock_dataset_utils,
|
||||
},
|
||||
):
|
||||
with patch.object(loader, "_setup_modelopt_quantization") as mock_setup:
|
||||
# Mock the _setup_modelopt_quantization to simulate checkpoint save
|
||||
def mock_setup_quantization(
|
||||
model,
|
||||
tokenizer,
|
||||
quant_cfg,
|
||||
quantized_ckpt_save_path=None,
|
||||
**kwargs,
|
||||
):
|
||||
# Simulate calibration and quantization
|
||||
mock_mtq.quantize(model, quant_cfg, forward_loop=MagicMock())
|
||||
mock_mtq.print_quant_summary(model)
|
||||
|
||||
# Save checkpoint if path provided
|
||||
if quantized_ckpt_save_path:
|
||||
mock_mto.save(model, quantized_ckpt_save_path)
|
||||
print(
|
||||
f"Quantized model saved to {quantized_ckpt_save_path}"
|
||||
)
|
||||
|
||||
mock_setup.side_effect = mock_setup_quantization
|
||||
|
||||
# Execute the load_model method
|
||||
result_model = loader.load_model(
|
||||
model_config=config_with_save, device_config=self.device_config
|
||||
)
|
||||
|
||||
# Verify the setup was called with save path
|
||||
mock_setup.assert_called_once()
|
||||
call_args = mock_setup.call_args
|
||||
# Check that the save path was passed correctly
|
||||
self.assertIn("quantized_ckpt_save_path", call_args[1])
|
||||
self.assertEqual(
|
||||
call_args[1]["quantized_ckpt_save_path"],
|
||||
"/path/to/save/checkpoint",
|
||||
)
|
||||
|
||||
# Verify save was called
|
||||
mock_mto.save.assert_called_once_with(
|
||||
self.mock_base_model, "/path/to/save/checkpoint"
|
||||
)
|
||||
|
||||
# Verify we get the expected model back
|
||||
self.assertEqual(result_model, self.mock_base_model)
|
||||
|
||||
def test_unified_quantization_flag_support(self):
|
||||
"""Test that ModelOptModelLoader supports unified quantization flags."""
|
||||
# Test modelopt_fp8
|
||||
config_fp8 = ModelConfig(
|
||||
model_path=self.model_path, quantization="modelopt_fp8"
|
||||
)
|
||||
self.assertEqual(config_fp8._get_modelopt_quant_type(), "fp8")
|
||||
|
||||
# Test modelopt_fp4
|
||||
config_fp4 = ModelConfig(
|
||||
model_path=self.model_path, quantization="modelopt_fp4"
|
||||
)
|
||||
self.assertEqual(config_fp4._get_modelopt_quant_type(), "nvfp4")
|
||||
|
||||
# Test auto-detection
|
||||
config_auto = ModelConfig(model_path=self.model_path, quantization="modelopt")
|
||||
# Should default to fp8 when no config is detected
|
||||
self.assertEqual(config_auto._get_modelopt_quant_type(), "fp8")
|
||||
|
||||
|
||||
class TestModelOptLoaderIntegration(CustomTestCase):
|
||||
"""Integration tests for ModelOptModelLoader with Engine API."""
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.get_model_loader")
|
||||
@patch("sglang.srt.entrypoints.engine.Engine.__init__")
|
||||
def test_engine_with_modelopt_quant_parameter(
|
||||
self, mock_engine_init, mock_get_model_loader
|
||||
):
|
||||
"""Test that Engine properly handles modelopt_quant parameter."""
|
||||
|
||||
# Mock the Engine.__init__ to avoid actual initialization
|
||||
mock_engine_init.return_value = None
|
||||
|
||||
# Mock get_model_loader to return our ModelOptModelLoader
|
||||
mock_loader = MagicMock(spec=ModelOptModelLoader)
|
||||
mock_get_model_loader.return_value = mock_loader
|
||||
|
||||
# Import here to avoid circular imports during test discovery
|
||||
# import sglang as sgl # Commented out since not directly used
|
||||
|
||||
# Test that we can create an engine with modelopt_quant parameter
|
||||
# This would normally trigger the ModelOptModelLoader selection
|
||||
try:
|
||||
engine_args = {
|
||||
"model_path": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
"modelopt_quant": "fp8",
|
||||
"log_level": "error", # Suppress logs during testing
|
||||
}
|
||||
|
||||
# This tests the parameter parsing and server args creation
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
server_args = ServerArgs(**engine_args)
|
||||
|
||||
# Verify that modelopt_quant is properly set
|
||||
self.assertEqual(server_args.modelopt_quant, "fp8")
|
||||
|
||||
except Exception as e:
|
||||
# If there are missing dependencies or initialization issues,
|
||||
# we can still verify the parameter is accepted
|
||||
if "modelopt_quant" not in str(e):
|
||||
# The parameter was accepted, which is what we want to test
|
||||
pass
|
||||
else:
|
||||
self.fail(f"modelopt_quant parameter not properly handled: {e}")
|
||||
|
||||
@patch("sglang.srt.model_loader.loader.get_model_loader")
|
||||
@patch("sglang.srt.entrypoints.engine.Engine.__init__")
|
||||
def test_engine_with_modelopt_quant_cli_argument(
|
||||
self, mock_engine_init, mock_get_model_loader
|
||||
):
|
||||
"""Test that CLI argument --modelopt-quant is properly parsed."""
|
||||
|
||||
# Mock the Engine.__init__ to avoid actual initialization
|
||||
mock_engine_init.return_value = None
|
||||
|
||||
# Mock get_model_loader to return our ModelOptModelLoader
|
||||
mock_loader = MagicMock(spec=ModelOptModelLoader)
|
||||
mock_get_model_loader.return_value = mock_loader
|
||||
|
||||
# Test CLI argument parsing
|
||||
import argparse
|
||||
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
# Create parser and add arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
|
||||
# Test parsing with modelopt_quant argument
|
||||
args = parser.parse_args(
|
||||
[
|
||||
"--model-path",
|
||||
"TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
"--modelopt-quant",
|
||||
"fp8",
|
||||
]
|
||||
)
|
||||
|
||||
# Convert to ServerArgs using the proper from_cli_args method
|
||||
server_args = ServerArgs.from_cli_args(args)
|
||||
|
||||
# Verify that modelopt_quant was properly parsed
|
||||
self.assertEqual(server_args.modelopt_quant, "fp8")
|
||||
self.assertEqual(server_args.model_path, "TinyLlama/TinyLlama-1.1B-Chat-v1.0")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,173 @@
|
||||
import asyncio
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed.device_mesh import init_device_mesh
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
from sglang.srt.weight_sync.utils import update_weights
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
register_cuda_ci(est_time=29, suite="stage-b-test-small-1-gpu")
|
||||
|
||||
|
||||
class AsyncEngine(Engine):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
async def update_weights_from_tensor(self, update_weights_request):
|
||||
return await self.tokenizer_manager.update_weights_from_tensor(
|
||||
update_weights_request, None
|
||||
)
|
||||
|
||||
|
||||
def is_distributed_available():
|
||||
"""Check if distributed training environment is available"""
|
||||
required_vars = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT"]
|
||||
return all(var in os.environ for var in required_vars)
|
||||
|
||||
|
||||
def setup_single_process_distributed():
|
||||
"""Setup distributed environment for single process testing"""
|
||||
if not is_distributed_available():
|
||||
os.environ["RANK"] = "0"
|
||||
os.environ["WORLD_SIZE"] = "1"
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = "12356"
|
||||
os.environ["LOCAL_RANK"] = "0"
|
||||
|
||||
|
||||
class TestUtilsUpdateWeights(unittest.TestCase):
|
||||
"""Test class for utils.update_weights function"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Setup distributed environment and test fixtures for the entire test class"""
|
||||
cls.setup_distributed()
|
||||
cls.setup_test_engine()
|
||||
cls.setup_test_model()
|
||||
cls.setup_device_mesh()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Cleanup after all tests"""
|
||||
if hasattr(cls, "engine") and cls.engine:
|
||||
cls.engine.shutdown()
|
||||
|
||||
# Cleanup distributed
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
@classmethod
|
||||
def setup_distributed(cls):
|
||||
"""Setup distributed environment for testing"""
|
||||
setup_single_process_distributed()
|
||||
|
||||
if not dist.is_initialized():
|
||||
try:
|
||||
dist.init_process_group(
|
||||
backend="nccl" if torch.cuda.is_available() else "gloo"
|
||||
)
|
||||
except Exception as e:
|
||||
raise unittest.SkipTest(
|
||||
f"Could not initialize distributed backend: {e}"
|
||||
)
|
||||
|
||||
cls.rank = dist.get_rank()
|
||||
cls.world_size = dist.get_world_size()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(cls.rank % torch.cuda.device_count())
|
||||
|
||||
# Set up environment variables
|
||||
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
||||
os.environ["NCCL_CUMEM_ENABLE"] = "0"
|
||||
os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "4"
|
||||
os.environ["CUDA_MODULE_LOADING"] = "AUTO"
|
||||
|
||||
@classmethod
|
||||
def setup_test_engine(cls):
|
||||
"""Setup test engine"""
|
||||
if cls.rank == 0:
|
||||
cls.engine = AsyncEngine(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
dtype="bfloat16",
|
||||
mem_fraction_static=0.3,
|
||||
enable_memory_saver=True,
|
||||
tp_size=cls.world_size,
|
||||
disable_cuda_graph=False,
|
||||
)
|
||||
else:
|
||||
cls.engine = None
|
||||
|
||||
@classmethod
|
||||
def setup_test_model(cls):
|
||||
"""Load test model"""
|
||||
try:
|
||||
cls.model = AutoModelForCausalLM.from_pretrained(
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
device_map="cpu",
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=(
|
||||
torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
raise unittest.SkipTest(f"Could not load test model: {e}")
|
||||
|
||||
@classmethod
|
||||
def setup_device_mesh(cls):
|
||||
"""Create device mesh for testing"""
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA not available for device mesh")
|
||||
|
||||
cls.device_mesh_key = "tp"
|
||||
cls.mesh = init_device_mesh(
|
||||
"cuda", (cls.world_size,), mesh_dim_names=(cls.device_mesh_key,)
|
||||
)
|
||||
|
||||
def create_test_params_batch(self, model, num_params=64):
|
||||
"""Create a batch of test parameters from the model"""
|
||||
param_names = []
|
||||
test_tensors = []
|
||||
|
||||
# Get first few parameters from the model for testing
|
||||
for i, (name, tensor) in enumerate(model.named_parameters()):
|
||||
if i >= num_params:
|
||||
break
|
||||
param_names.append(name)
|
||||
# Create test tensor with known values, matching original shape and dtype
|
||||
test_tensor = torch.full_like(tensor, 1.5, dtype=tensor.dtype).cuda()
|
||||
test_tensors.append(test_tensor)
|
||||
|
||||
return list(zip(param_names, test_tensors))
|
||||
|
||||
def test_utils_update_weights(self):
|
||||
"""Test basic functionality of utils.update_weights"""
|
||||
|
||||
async def async_test():
|
||||
# Create test parameters batch
|
||||
params_batch = self.create_test_params_batch(self.model, num_params=2)
|
||||
|
||||
# Test the utils.update_weights function
|
||||
result = await update_weights(
|
||||
engine=self.engine,
|
||||
params_batch=params_batch,
|
||||
device_mesh_key=self.device_mesh_key,
|
||||
device_mesh=self.mesh,
|
||||
load_format=None,
|
||||
)
|
||||
|
||||
self.assertIn("Success", result)
|
||||
|
||||
# Run the async test
|
||||
asyncio.run(async_test())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user