912 lines
36 KiB
Python
912 lines
36 KiB
Python
"""
|
|
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 types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
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.linear import ReplicatedLinear
|
|
from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method
|
|
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
|
|
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
|
|
from sglang.srt.layers.quantization.modelopt_quant import (
|
|
ModelOptFp4Config,
|
|
ModelOptFp4LinearMethod,
|
|
ModelOptFp8Config,
|
|
ModelOptMixedPrecisionConfig,
|
|
ModelOptNvFp4A16LinearMethod,
|
|
)
|
|
from sglang.srt.model_loader.loader import (
|
|
DefaultModelLoader,
|
|
ModelOptModelLoader,
|
|
get_model_loader,
|
|
)
|
|
from sglang.srt.model_loader.weight_utils import get_quant_config
|
|
from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
|
|
from sglang.srt.models.utils import WeightsMapper
|
|
from sglang.srt.utils import get_device
|
|
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, stage="base-b", runner_config="1-gpu-small")
|
|
|
|
|
|
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=get_device())
|
|
|
|
# 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.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_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")
|
|
|
|
|
|
class TestParseQuantHfConfig(CustomTestCase):
|
|
"""Tests for _parse_quant_hf_config and _parse_modelopt_quant_config.
|
|
|
|
Regression tests for the fix where quant_method='modelopt' ignoring quant_algo.
|
|
"""
|
|
|
|
# (quant_config_input, expected_quant_method)
|
|
_MODELOPT_CASES = [
|
|
({"quant_method": "modelopt", "quant_algo": "FP8"}, "modelopt_fp8"),
|
|
({"quant_method": "modelopt", "quant_algo": "FP4"}, "modelopt_fp4"),
|
|
({"quant_method": "modelopt", "quant_algo": "NVFP4"}, "modelopt_fp4"),
|
|
({"quant_algo": "NVFP4_AWQ"}, "modelopt_fp4"),
|
|
({"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION"}, "w4afp8"),
|
|
({"quant_algo": "FP8"}, "modelopt_fp8"),
|
|
({"quant_algo": "MXFP8"}, "mxfp8"),
|
|
({"quant_algo": "FP4"}, "modelopt_fp4"),
|
|
({"quant_algo": "MIXED_PRECISION"}, "w4afp8"),
|
|
({"quant_method": "modelopt"}, "modelopt"),
|
|
]
|
|
|
|
def setUp(self):
|
|
"""Set up a real ModelConfig using TinyLlama (already used elsewhere)."""
|
|
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_mp_is_initialized = patch(
|
|
"sglang.srt.distributed.parallel_state.model_parallel_is_initialized",
|
|
return_value=True,
|
|
)
|
|
self.mock_mp_is_initialized.start()
|
|
|
|
self.model_config = ModelConfig(
|
|
model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.mock_tp_rank.stop()
|
|
self.mock_mp_is_initialized.stop()
|
|
|
|
def test_modelopt_quant_parsing(self):
|
|
"""Modelopt quant configs must resolve to the correct quant_method."""
|
|
for quant_cfg_input, expected in self._MODELOPT_CASES:
|
|
with self.subTest(quant_cfg=quant_cfg_input):
|
|
self.model_config.hf_config.quantization_config = dict(quant_cfg_input)
|
|
result = self.model_config._parse_quant_hf_config()
|
|
self.assertEqual(result["quant_method"], expected)
|
|
|
|
def test_awq_flat_config_defaults_group_size(self):
|
|
"""NVFP4_AWQ flat config.json omits group_size; from_config must default it to 16."""
|
|
cfg = ModelOptFp4Config.from_config(
|
|
{
|
|
"quant_algo": "NVFP4_AWQ",
|
|
"ignore": ["lm_head"],
|
|
"quant_method": "modelopt",
|
|
}
|
|
)
|
|
self.assertEqual(cfg.group_size, 16)
|
|
self.assertTrue(cfg.is_awq)
|
|
|
|
def test_modelopt_mxfp8_config(self):
|
|
"""ModelOpt MXFP8 metadata must select block scales and retain FP8 KV policy."""
|
|
model_config = ModelConfig.__new__(ModelConfig)
|
|
for kv_cache_config in (
|
|
{"kv_cache_quant_algo": "FP8"},
|
|
{"kv_cache_scheme": {"type": "float", "num_bits": 8}},
|
|
):
|
|
with self.subTest(kv_cache_config=kv_cache_config):
|
|
result = model_config._parse_modelopt_quant_config(
|
|
{
|
|
"quantization": {
|
|
"quant_algo": "MXFP8",
|
|
"group_size": 32,
|
|
"exclude_modules": ["lm_head"],
|
|
**kv_cache_config,
|
|
}
|
|
}
|
|
)
|
|
self.assertEqual(result["quant_method"], "mxfp8")
|
|
self.assertEqual(result["scale_fmt"], "ue8m0")
|
|
|
|
quant_config = Fp8Config.from_config(result)
|
|
self.assertEqual(quant_config.get_name(), "mxfp8")
|
|
self.assertEqual(quant_config.activation_scheme, "dynamic")
|
|
self.assertEqual(quant_config.weight_block_size, [1, 32])
|
|
self.assertIn("lm_head", quant_config.ignored_layers)
|
|
self.assertEqual(quant_config.kv_cache_quant_algo, "FP8")
|
|
|
|
def test_modelopt_mxfp8_override(self):
|
|
"""Generic ModelOpt selection must not route MXFP8 to scalar FP8."""
|
|
self.assertEqual(
|
|
ModelOptFp8Config.override_quantization_method(
|
|
{"quant_algo": "MXFP8"}, "modelopt"
|
|
),
|
|
"mxfp8",
|
|
)
|
|
|
|
def test_modelopt_mxfp8_weight_loading(self):
|
|
"""ModelOpt MXFP8 block scales must reach native scale parameters."""
|
|
weight = torch.empty(1)
|
|
weights = [
|
|
("model.q_proj.weight_scale", weight),
|
|
("model.q_proj.input_weight_scale", weight),
|
|
("model.q_proj.weight_scale_inv", weight),
|
|
]
|
|
|
|
def load_names(quant_config):
|
|
model = nn.Module()
|
|
model.quant_config = quant_config
|
|
loaded_names = []
|
|
model.load_weights = lambda weights: loaded_names.extend(
|
|
name for name, _ in weights
|
|
)
|
|
with patch(
|
|
"sglang.srt.model_loader.loader.is_cuda_alike", return_value=False
|
|
):
|
|
DefaultModelLoader.load_weights_and_postprocess(
|
|
model, iter(weights), torch.device("cpu")
|
|
)
|
|
return loaded_names
|
|
|
|
mxfp8_config = Fp8Config(
|
|
is_checkpoint_fp8_serialized=True,
|
|
activation_scheme="dynamic",
|
|
weight_block_size=[1, 32],
|
|
use_mxfp8=True,
|
|
)
|
|
self.assertEqual(
|
|
load_names(mxfp8_config),
|
|
[
|
|
"model.q_proj.weight_scale_inv",
|
|
"model.q_proj.input_weight_scale",
|
|
"model.q_proj.weight_scale_inv",
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
load_names(Fp8Config(is_checkpoint_fp8_serialized=True)),
|
|
[
|
|
"model.q_proj.weight_scale",
|
|
"model.q_proj.input_weight_scale",
|
|
"model.q_proj.weight_scale_inv",
|
|
],
|
|
)
|
|
|
|
def test_non_modelopt_quant_method_unchanged(self):
|
|
"""Non-modelopt quant_method (e.g. 'gptq') must NOT enter the modelopt path."""
|
|
self.model_config.hf_config.quantization_config = {
|
|
"quant_method": "gptq",
|
|
"bits": 4,
|
|
}
|
|
result = self.model_config._parse_quant_hf_config()
|
|
self.assertEqual(result["quant_method"], "gptq")
|
|
self.assertNotIn("quant_algo", result)
|
|
|
|
def test_inherited_draft_modelopt_fp4_accepts_fp8_checkpoint(self):
|
|
# ServerArgs has already copied the target's modelopt_fp4 request to the
|
|
# draft. Compatible FP8 metadata must not replace it with plain fp8.
|
|
self.model_config.quantization = "modelopt_fp4"
|
|
self.model_config.is_draft_model = True
|
|
self.model_config.is_draft_quantization_explicit = False
|
|
with (
|
|
patch.object(
|
|
self.model_config,
|
|
"_parse_quant_hf_config",
|
|
return_value={"quant_method": "fp8"},
|
|
),
|
|
patch.object(
|
|
self.model_config,
|
|
"_find_quant_modelslim_config",
|
|
return_value=None,
|
|
),
|
|
):
|
|
self.model_config._verify_quantization()
|
|
|
|
# Keeping modelopt_fp4 selects online FP8-to-NVFP4 conversion for
|
|
# eligible MoE experts; this test stops at quantization-method routing.
|
|
self.assertEqual(self.model_config.quantization, "modelopt_fp4")
|
|
|
|
|
|
class TestModelOptFp4LoaderSelection(CustomTestCase):
|
|
def test_draft_modelopt_fp4_uses_checkpoint_exclusions(self):
|
|
cases = (
|
|
# Excluded MTP experts are unpacked, so an explicit draft request
|
|
# replaces the serialized config with online weight quantization.
|
|
("explicit embedded draft", True, ["mtp.layers.0*"], False),
|
|
# MTP experts present in the serialized checkpoint stay serialized.
|
|
("explicit serialized draft", True, [], True),
|
|
# Inherited target quantization does not override draft exclusions.
|
|
("inherited embedded draft", False, ["mtp.layers.0*"], True),
|
|
)
|
|
for name, is_explicit, ignored_layers, is_serialized in cases:
|
|
with self.subTest(name=name):
|
|
model_config = SimpleNamespace(
|
|
model_path="target-model",
|
|
quantization="modelopt_fp4",
|
|
is_draft_model=True,
|
|
is_draft_quantization_explicit=is_explicit,
|
|
hf_config=SimpleNamespace(
|
|
quantization_config={
|
|
"quant_algo": "NVFP4",
|
|
"group_size": 16,
|
|
"ignore": ignored_layers,
|
|
}
|
|
),
|
|
)
|
|
|
|
config = get_quant_config(model_config, LoadConfig(), {})
|
|
|
|
self.assertEqual(config.get_name(), "modelopt_fp4")
|
|
self.assertEqual(config.is_checkpoint_nvfp4_serialized, is_serialized)
|
|
|
|
def test_unquantized_modelopt_fp4_preserves_modelopt_workflows(self):
|
|
model_config = SimpleNamespace(
|
|
quantization="modelopt_fp4",
|
|
_is_already_quantized=lambda: False,
|
|
)
|
|
|
|
# Online conversion runs through the regular per-layer weight loaders.
|
|
online_loader = get_model_loader(LoadConfig(), model_config)
|
|
self.assertIsInstance(online_loader, DefaultModelLoader)
|
|
self.assertNotIsInstance(online_loader, ModelOptModelLoader)
|
|
|
|
# Explicit ModelOpt checkpoint/export workflows still need its loader.
|
|
for option in (
|
|
"modelopt_checkpoint_restore_path",
|
|
"modelopt_checkpoint_save_path",
|
|
"modelopt_export_path",
|
|
):
|
|
with self.subTest(option=option):
|
|
loader = get_model_loader(
|
|
LoadConfig(**{option: "/tmp/modelopt"}), model_config
|
|
)
|
|
self.assertIsInstance(loader, ModelOptModelLoader)
|
|
|
|
|
|
class TestModelOptMixedPrecisionConfig(CustomTestCase):
|
|
def test_minimax_mixed_precision_resolves_runtime_names_and_mxfp8(self):
|
|
quant_config = ModelOptMixedPrecisionConfig.from_config(
|
|
{
|
|
"quant_algo": "MIXED_PRECISION",
|
|
"weight_block_size": [1, 32],
|
|
"exclude_modules": ["language_model.lm_head"],
|
|
"quantized_layers": {
|
|
"language_model.model.layers.3.self_attn.q_proj": {
|
|
"quant_algo": "MXFP8"
|
|
},
|
|
"language_model.model.layers.3.self_attn.k_proj": {
|
|
"quant_algo": "MXFP8"
|
|
},
|
|
"language_model.model.layers.3.self_attn.v_proj": {
|
|
"quant_algo": "MXFP8"
|
|
},
|
|
"language_model.model.layers.3.block_sparse_moe.experts.0.w1": {
|
|
"quant_algo": "NVFP4",
|
|
"group_size": 16,
|
|
},
|
|
"language_model.model.layers.3.block_sparse_moe.shared_experts.gate_proj": {
|
|
"quant_algo": "MXFP8"
|
|
},
|
|
},
|
|
"packed_modules_mapping": {
|
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
|
"gate_up_proj": ["gate_proj", "up_proj"],
|
|
},
|
|
}
|
|
)
|
|
quant_config.apply_weight_name_mapper(
|
|
MiniMaxM3SparseForCausalLM.hf_to_sglang_mapper
|
|
)
|
|
|
|
self.assertEqual(
|
|
quant_config._resolve_quant_algo(
|
|
"language_model.model.layers.3.mlp.experts"
|
|
),
|
|
"NVFP4",
|
|
)
|
|
self.assertEqual(
|
|
quant_config._resolve_quant_algo(
|
|
"language_model.model.layers.3.mlp.shared_experts.gate_up_proj"
|
|
),
|
|
"MXFP8",
|
|
)
|
|
|
|
# Type dispatch only needs a LinearBase instance; skip GPU weight setup.
|
|
linear = ReplicatedLinear.__new__(ReplicatedLinear)
|
|
method = quant_config.get_quant_method(
|
|
linear, "language_model.model.layers.3.self_attn.qkv_proj"
|
|
)
|
|
self.assertIsInstance(method, Fp8LinearMethod)
|
|
self.assertTrue(method.use_mxfp8)
|
|
self.assertEqual(quant_config.mxfp8_config.weight_block_size, [1, 32])
|
|
self.assertEqual(
|
|
quant_config.exclude_modules,
|
|
["language_model.lm_head", "lm_head"],
|
|
)
|
|
|
|
def test_nemotron_mixed_precision_with_nvfp4_layers_uses_modelopt_mixed(self):
|
|
model_config = ModelConfig.__new__(ModelConfig)
|
|
model_config.hf_config = MagicMock()
|
|
model_config.hf_config.model_type = "nemotron_h"
|
|
model_config.hf_config.architectures = ["NemotronHForCausalLM"]
|
|
|
|
result = model_config._parse_modelopt_quant_config(
|
|
{
|
|
"quantization": {
|
|
"quant_algo": "MIXED_PRECISION",
|
|
"quantized_layers": {
|
|
"backbone.layers.0.mixer.in_proj": {"quant_algo": "FP8"},
|
|
"backbone.layers.0.mixer.out_proj": {"quant_algo": "FP8"},
|
|
"backbone.layers.1.mixer.experts.0.up_proj": {
|
|
"quant_algo": "NVFP4",
|
|
"group_size": 16,
|
|
},
|
|
"backbone.layers.1.mixer.experts.0.down_proj": {
|
|
"quant_algo": "NVFP4",
|
|
"group_size": 16,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
)
|
|
|
|
self.assertEqual(result["quant_method"], "modelopt_mixed")
|
|
|
|
def test_qwen_mixed_precision_with_nvfp4a16_layers_uses_modelopt_mixed(self):
|
|
model_config = ModelConfig.__new__(ModelConfig)
|
|
model_config.hf_config = MagicMock()
|
|
model_config.hf_config.model_type = "qwen3_5_moe"
|
|
model_config.hf_config.architectures = ["Qwen3_5MoeForConditionalGeneration"]
|
|
|
|
result = model_config._parse_modelopt_quant_config(
|
|
{
|
|
"quantization": {
|
|
"quant_algo": "MIXED_PRECISION",
|
|
"quantized_layers": {
|
|
"lm_head": {"quant_algo": "W4A16_NVFP4", "group_size": 16},
|
|
"model.language_model.layers.0.mlp.shared_expert.up_proj": {
|
|
"quant_algo": "W4A16_NVFP4",
|
|
"group_size": 16,
|
|
},
|
|
"model.language_model.layers.0.linear_attn.in_proj_qkv": {
|
|
"quant_algo": "FP8"
|
|
},
|
|
},
|
|
}
|
|
}
|
|
)
|
|
|
|
self.assertEqual(result["quant_method"], "modelopt_mixed")
|
|
|
|
def test_mixed_precision_override_does_not_hijack_w4afp8(self):
|
|
self.assertIsNone(
|
|
ModelOptMixedPrecisionConfig.override_quantization_method(
|
|
{"quant_method": "w4afp8", "quant_algo": "MIXED_PRECISION"},
|
|
"w4afp8",
|
|
)
|
|
)
|
|
|
|
@patch(
|
|
"sglang.srt.layers.quantization.modelopt_quant.envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get",
|
|
return_value=True,
|
|
)
|
|
def test_explicit_nvfp4_per_token_activation_false_overrides_env(self, _):
|
|
config = ModelOptFp4Config(
|
|
is_checkpoint_nvfp4_serialized=True,
|
|
group_size=16,
|
|
use_per_token_activation=False,
|
|
)
|
|
|
|
self.assertFalse(config.use_per_token_activation)
|
|
|
|
def test_lm_head_guard_accepts_modelopt_fp4_marlin_runtime_state(self):
|
|
lm_head = nn.Module()
|
|
lm_head.weight = nn.Parameter(
|
|
torch.empty(128, 496640, dtype=torch.int32), requires_grad=False
|
|
)
|
|
lm_head.weight_scale = nn.Parameter(torch.empty(1))
|
|
lm_head.weight_global_scale = nn.Parameter(torch.empty(1))
|
|
lm_head.workspace = torch.empty(1)
|
|
lm_head.input_size_per_partition = 2048
|
|
lm_head.output_size_per_partition = 128000
|
|
|
|
self.assertTrue(
|
|
should_apply_lm_head_quant_method(
|
|
lm_head, ModelOptNvFp4A16LinearMethod(ModelOptFp4Config())
|
|
)
|
|
)
|
|
|
|
def test_lm_head_guard_rejects_stale_modelopt_fp4_method_on_dense_head(self):
|
|
lm_head = nn.Module()
|
|
lm_head.weight = nn.Parameter(torch.empty(128000, 2048))
|
|
|
|
self.assertFalse(
|
|
should_apply_lm_head_quant_method(
|
|
lm_head, ModelOptFp4LinearMethod(ModelOptFp4Config())
|
|
)
|
|
)
|
|
|
|
def test_lm_head_guard_rejects_stale_modelopt_fp4_attrs_on_dense_head(self):
|
|
lm_head = nn.Module()
|
|
lm_head.weight = nn.Parameter(torch.empty(128000, 2048))
|
|
lm_head.weight_scale = nn.Parameter(torch.empty(1))
|
|
lm_head.weight_global_scale = nn.Parameter(torch.empty(1))
|
|
lm_head.workspace = torch.empty(1)
|
|
lm_head.input_size_per_partition = 2048
|
|
lm_head.output_size_per_partition = 128000
|
|
|
|
self.assertFalse(
|
|
should_apply_lm_head_quant_method(
|
|
lm_head, ModelOptNvFp4A16LinearMethod(ModelOptFp4Config())
|
|
)
|
|
)
|
|
|
|
def test_mixed_precision_quant_layer_resolution_after_mapping(self):
|
|
quant_config = ModelOptMixedPrecisionConfig.from_config(
|
|
{
|
|
"quant_algo": "MIXED_PRECISION",
|
|
"quantized_layers": {
|
|
"backbone.layers.0.mixer.in_proj": {"quant_algo": "FP8"},
|
|
"backbone.layers.1.mixer.experts.0.up_proj": {
|
|
"quant_algo": "NVFP4",
|
|
"group_size": 16,
|
|
},
|
|
"backbone.layers.2.mixer.q_proj": {"quant_algo": "FP8"},
|
|
"backbone.layers.2.mixer.k_proj": {"quant_algo": "FP8"},
|
|
"backbone.layers.2.mixer.v_proj": {"quant_algo": "FP8"},
|
|
},
|
|
"packed_modules_mapping": {
|
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
|
},
|
|
}
|
|
)
|
|
quant_config.apply_weight_name_mapper(
|
|
WeightsMapper(orig_to_new_prefix={"backbone.": "model."})
|
|
)
|
|
|
|
self.assertEqual(
|
|
quant_config._resolve_quant_algo("model.layers.0.mixer.in_proj"),
|
|
"FP8",
|
|
)
|
|
self.assertEqual(
|
|
quant_config._resolve_quant_algo("model.layers.1.mixer.experts"),
|
|
"NVFP4",
|
|
)
|
|
self.assertEqual(
|
|
quant_config._resolve_quant_algo("model.layers.2.mixer.qkv_proj"),
|
|
"FP8",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|