Fix external_models import path and migrate model loading tests (#16458)

This commit is contained in:
Alison Shao
2026-01-08 23:43:49 -08:00
committed by GitHub
parent 9d4d57dbfa
commit e46f79431b
8 changed files with 24 additions and 21 deletions
@@ -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()