Support ModelOpt MXFP8 checkpoints (#32538)
This commit is contained in:
@@ -1249,8 +1249,34 @@ class ModelConfig:
|
||||
return {"quant_method": "w4afp8", "quant_algo": quant_algo}
|
||||
elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo):
|
||||
return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo}
|
||||
elif quant_algo and "FP8" in quant_algo:
|
||||
elif quant_algo == "FP8":
|
||||
return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo}
|
||||
elif quant_algo == "MXFP8":
|
||||
group_size = json_quant_configs.get("group_size", 32)
|
||||
ignored_layers = json_quant_configs.get(
|
||||
"exclude_modules", json_quant_configs.get("ignore")
|
||||
)
|
||||
kv_cache_quant_algo = json_quant_configs.get("kv_cache_quant_algo")
|
||||
if kv_cache_quant_algo is None:
|
||||
kv_cache_scheme = json_quant_configs.get("kv_cache_scheme")
|
||||
if (
|
||||
isinstance(kv_cache_scheme, dict)
|
||||
and kv_cache_scheme.get("type") == "float"
|
||||
and kv_cache_scheme.get("num_bits") == 8
|
||||
):
|
||||
kv_cache_quant_algo = "FP8"
|
||||
parsed = {
|
||||
"quant_method": "mxfp8",
|
||||
"quant_algo": quant_algo,
|
||||
"activation_scheme": "dynamic",
|
||||
"weight_block_size": [1, group_size],
|
||||
"scale_fmt": "ue8m0",
|
||||
}
|
||||
if ignored_layers is not None:
|
||||
parsed["modules_to_not_convert"] = ignored_layers
|
||||
if kv_cache_quant_algo is not None:
|
||||
parsed["kv_cache_quant_algo"] = kv_cache_quant_algo
|
||||
return parsed
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
@@ -190,7 +190,9 @@ class QuantizationConfig(ABC):
|
||||
|
||||
# If user specified generic "modelopt", auto-detect the specific method
|
||||
if user_quant == "modelopt":
|
||||
if "FP8" in quant_algo:
|
||||
if quant_algo == "MXFP8":
|
||||
return "mxfp8"
|
||||
elif quant_algo == "FP8":
|
||||
return "modelopt_fp8"
|
||||
elif "NVFP4" in quant_algo or "FP4" in quant_algo:
|
||||
return "modelopt_fp4"
|
||||
|
||||
@@ -235,6 +235,7 @@ class Fp8Config(QuantizationConfig):
|
||||
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
|
||||
use_mxfp8: bool = False,
|
||||
is_fp4_experts: bool = False,
|
||||
kv_cache_quant_algo: Optional[str] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# DSV4 mxfp4-packed (True) vs converted FP8 (False); injected by
|
||||
@@ -258,6 +259,7 @@ class Fp8Config(QuantizationConfig):
|
||||
)
|
||||
self.packed_modules_mapping = packed_modules_mapping or {}
|
||||
self.use_mxfp8 = use_mxfp8
|
||||
self.kv_cache_quant_algo = kv_cache_quant_algo
|
||||
if weight_block_size is not None:
|
||||
if not is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
@@ -322,6 +324,9 @@ class Fp8Config(QuantizationConfig):
|
||||
normalized.append(f"model.{base}")
|
||||
ignored_layers = normalized
|
||||
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
|
||||
kv_cache_quant_algo = cls.get_from_keys_or(
|
||||
config, ["kv_cache_quant_algo"], None
|
||||
)
|
||||
if use_mxfp8:
|
||||
# MXFP8 (OCP) spec fixes block size to [1, 32]; ckpt field is metadata only.
|
||||
if weight_block_size is not None and weight_block_size != [1, 32]:
|
||||
@@ -337,6 +342,7 @@ class Fp8Config(QuantizationConfig):
|
||||
weight_block_size=weight_block_size,
|
||||
packed_modules_mapping=packed_modules_mapping,
|
||||
use_mxfp8=use_mxfp8,
|
||||
kv_cache_quant_algo=kv_cache_quant_algo,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
|
||||
@@ -465,11 +465,11 @@ class ModelOptFp8Config(ModelOptQuantConfig):
|
||||
raise ValueError(
|
||||
"Cannot find 'quant_algo' in the model's quantization config. "
|
||||
)
|
||||
if "FP8" not in quant_method:
|
||||
if quant_method != "FP8":
|
||||
raise ValueError(
|
||||
"ModelOptFp8Config only supports static FP8 quantization in SGLang. "
|
||||
"For FP4 quantization, use ModelOptFp4Config. "
|
||||
"Check the quantization config for your model's configuration."
|
||||
"ModelOptFp8Config only supports regular FP8 quantization, "
|
||||
f"but found {quant_method!r}. Use the native 'mxfp8' "
|
||||
"quantization method for MXFP8 or ModelOptFp4Config for FP4."
|
||||
)
|
||||
|
||||
return cls(
|
||||
|
||||
@@ -835,6 +835,15 @@ class DefaultModelLoader(BaseModelLoader):
|
||||
|
||||
quant_config = getattr(model, "quant_config", None)
|
||||
is_nvfp4_online = getattr(quant_config, "is_nvfp4_online", False)
|
||||
is_mxfp8 = quant_config is not None and quant_config.get_name() == "mxfp8"
|
||||
if is_mxfp8:
|
||||
weights = (
|
||||
(
|
||||
f"{name}_inv" if name.endswith(".weight_scale") else name,
|
||||
loaded_weight,
|
||||
)
|
||||
for name, loaded_weight in weights
|
||||
)
|
||||
|
||||
if is_nvfp4_online:
|
||||
# Scope exact FP4 quantization math to load-time conversion only;
|
||||
|
||||
@@ -17,14 +17,15 @@ 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 Fp8LinearMethod
|
||||
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 ModelOptModelLoader
|
||||
from sglang.srt.model_loader.loader import DefaultModelLoader, ModelOptModelLoader
|
||||
from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
|
||||
from sglang.srt.models.utils import WeightsMapper
|
||||
from sglang.srt.utils import get_device
|
||||
@@ -462,6 +463,7 @@ class TestParseQuantHfConfig(CustomTestCase):
|
||||
({"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"),
|
||||
@@ -509,6 +511,90 @@ class TestParseQuantHfConfig(CustomTestCase):
|
||||
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 = {
|
||||
|
||||
Reference in New Issue
Block a user