fix: validate quantization methods before startup

lightllm/common/quantization/__init__.py

@@ -2,11 +2,6 @@
 import yaml
 import collections
 from .registry import QUANTMETHODS
-from .w8a8 import *
-from .w8a8gx import *
-from .deepgemm import *
-from .awq import *
-from .no_quant import *
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.device_utils import is_sm100_gpu
 
@@ -19,6 +14,42 @@
 SUPPORTED_EXPERT_DTYPES = tuple(EXPERT_DTYPE_TO_QUANT_TYPE)
 
 
+def load_quant_config(quant_type="none", custom_cfg_path=None, platform="cuda"):
+    if custom_cfg_path is None:
+        QUANTMETHODS.validate(quant_type, platform)
+        return None
+    try:
+        with open(custom_cfg_path, "r") as file:
+            data = yaml.safe_load(file)
+        if not isinstance(data, dict):
+            raise ValueError("quantization config must be a mapping")
+        if "quant_type" not in data:
+            raise ValueError("missing required field 'quant_type'")
+        QUANTMETHODS.validate(data["quant_type"], platform)
+        mix_bits = data.get("mix_bits", [])
+        if not isinstance(mix_bits, list):
+            raise ValueError("mix_bits must be a list")
+        for index, layer_quant_cfg in enumerate(mix_bits):
+            try:
+                if not isinstance(layer_quant_cfg, dict):
+                    raise ValueError("entry must be a mapping")
+                for key in ("name", "quant_type"):
+                    if key not in layer_quant_cfg:
+                        raise ValueError(f"missing required field '{key}'")
+                if not isinstance(layer_quant_cfg["name"], str) or not layer_quant_cfg["name"]:
+                    raise ValueError("name must be a non-empty string")
+                QUANTMETHODS.validate(layer_quant_cfg["quant_type"], platform)
+                if "layer_nums" in layer_quant_cfg:
+                    layer_nums = layer_quant_cfg["layer_nums"]
+                    if not isinstance(layer_nums, list) or any(type(n) is not int or n < 0 for n in layer_nums):
+                        raise ValueError("layer_nums must be a list of non-negative integers")
+            except ValueError as e:
+                raise ValueError(f"mix_bits[{index}]: {e}") from e
+    except (OSError, ValueError, yaml.YAMLError) as e:
+        raise ValueError(f"{custom_cfg_path}: {e}") from e
+    return data
+
+
 class Quantcfg:
     def __init__(self, network_config, quant_type="none", custom_cfg_path=None, expert_dtype=None):
         self.layer_num = network_config["n_layer"]
@@ -95,6 +126,8 @@ def _mapping_quant_method(self):
             self._mapping_expert_quant_method()
 
         elif self.hf_quantization_method == "awq":
+            from .awq import is_awq_marlin_compatible
+
             self.quant_type = "awq"
             if is_awq_marlin_compatible(self.hf_quantization_config):
                 self.quant_type = "awq_marlin"
@@ -105,12 +138,10 @@ def _mapping_quant_method(self):
 
     def _parse_custom_cfg(self, custom_cfg_path):
         self.quant_cfg = collections.defaultdict(dict)
-        if custom_cfg_path is None:
+        data = load_quant_config(self.quant_type, custom_cfg_path)
+        if data is None:
             return
 
-        with open(custom_cfg_path, "r") as file:
-            data = yaml.safe_load(file)
-
         self.quant_type = data["quant_type"]
         for layer_quant_cfg in data.get("mix_bits", []):
             name = layer_quant_cfg["name"]

lightllm/common/quantization/awq.py

@@ -2,7 +2,6 @@
 from typing import Any, Optional, Tuple, List
 
 from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
-from lightllm.common.quantization.registry import QUANTMETHODS
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.utils.log_utils import init_logger
 
@@ -66,7 +65,6 @@ def method_name(self):
         return "awq-base"
 
 
-@QUANTMETHODS.register("awq", platform="cuda")
 class AWQW4A16QuantizationMethod(AWQBaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -135,7 +133,6 @@ def _create_weight(
         return mm_param, mm_param_list
 
 
-@QUANTMETHODS.register("awq_marlin", platform="cuda")
 class AWQMARLINW4A16QuantizationMethod(AWQBaseQuantizationMethod):
     def __init__(self):
         super().__init__()

lightllm/common/quantization/deepgemm.py

@@ -3,7 +3,6 @@
 from typing import Optional, List, Union, Tuple
 
 from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
-from lightllm.common.quantization.registry import QUANTMETHODS
 from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
 from lightllm.utils.device_utils import is_sm100_gpu
 from lightllm.utils.log_utils import init_logger
@@ -45,7 +44,6 @@ def method_name(self):
         return "base-deepgemm"
 
 
-@QUANTMETHODS.register("fp8w8a8-b128-deepgemm", platform="cuda")
 class DeepGEMMFP8w8a8B128QuantizationMethod(DeepGEMMBaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -160,7 +158,6 @@ def _create_weight(
         return mm_param, mm_param_list
 
 
-@QUANTMETHODS.register("fp4fp8-b32-deepgemm", platform="cuda")
 class DeepGEMMFP8FP4B32QuantizationMethod(DeepGEMMBaseQuantizationMethod):
     def __init__(self):
         super().__init__()

lightllm/common/quantization/no_quant.py

@@ -2,11 +2,8 @@
 from typing import Optional, List, Union, Tuple
 
 from lightllm.common.quantization.quantize_method import QuantizationMethod, WeightPack
-from lightllm.common.quantization.registry import QUANTMETHODS
 
 
-@QUANTMETHODS.register("none", platform="musa")
-@QUANTMETHODS.register("none", platform="cuda")
 class NoQuantization(QuantizationMethod):
     """No quantization - uses full precision weights."""
 

lightllm/common/quantization/registry.py

@@ -1,4 +1,8 @@
-from .quantize_method import QuantizationMethod
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .quantize_method import QuantizationMethod
 
 
 class QuantMethodFactory:
@@ -18,15 +22,45 @@ def decorator(cls):
 
         return decorator
 
-    def get(self, key, platform="cuda", *args, **kwargs) -> "QuantizationMethod":
-        quant_method_class_dict = self._quant_methods.get(key)
-        if not quant_method_class_dict:
-            raise ValueError(f"QuantMethod '{key}' not supported.")
+    def supported_methods(self, platform="cuda"):
+        return tuple(
+            sorted(name for name, methods in self._quant_methods.items() if platform is None or platform in methods)
+        )
+
+    def validate(self, key, platform="cuda"):
+        if not isinstance(key, str) or platform not in self._quant_methods.get(key, {}):
+            raise ValueError(
+                f"QuantMethod '{key}' for platform '{platform}' not supported. "
+                f"Supported methods: {', '.join(self.supported_methods(platform))}"
+            )
 
-        quant_method_class = quant_method_class_dict.get(platform)
-        if quant_method_class is None:
-            raise ValueError(f"QuantMethod '{key}' for platform '{platform}' not supported.")
+    def get(self, key, platform="cuda", *args, **kwargs) -> "QuantizationMethod":
+        self.validate(key, platform)
+        quant_method_class = self._quant_methods[key][platform]
+        if isinstance(quant_method_class, str):
+            module, class_name = quant_method_class.rsplit(".", 1)
+            quant_method_class = getattr(import_module(f".{module}", __package__), class_name)
+            self._quant_methods[key][platform] = quant_method_class
         return quant_method_class()
 
 
 QUANTMETHODS = QuantMethodFactory()
+
+for names, method in {
+    ("w8a8-vllm", "w8a8"): "w8a8.w8a8QuantizationMethod",
+    ("fp8w8a8-vllm", "fp8w8a8"): "w8a8.FP8w8a8QuantizationMethod",
+    "fp8w8a8-pt-vllm": "w8a8.FP8w8a8PerTensorVllmQuantizationMethod",
+    "fp8w8a8-pt-sgl": "w8a8.FP8w8a8PerTensorSglQuantizationMethod",
+    ("fp8w8a8-pt-triton", "fp8w8a8-pt"): "w8a8.FP8w8a8PerTensorTritonQuantizationMethod",
+    "fp8w8a8-b128-vllm": "w8a8.FP8w8a8B128VllmQuantizationMethod",
+    ("fp8w8a8-b128-triton", "fp8w8a8-b128"): "w8a8.FP8w8a8B128TritonQuantizationMethod",
+    ("fp8w8a8g128-triton", "fp8w8a8g128"): "w8a8gx.FP8w8a8g128QuantizationMethod",
+    ("fp8w8a8g64-triton", "fp8w8a8g64"): "w8a8gx.FP8w8a8g64QuantizationMethod",
+    "fp8w8a8-b128-deepgemm": "deepgemm.DeepGEMMFP8w8a8B128QuantizationMethod",
+    "fp4fp8-b32-deepgemm": "deepgemm.DeepGEMMFP8FP4B32QuantizationMethod",
+    "awq": "awq.AWQW4A16QuantizationMethod",
+    "awq_marlin": "awq.AWQMARLINW4A16QuantizationMethod",
+    "none": "no_quant.NoQuantization",
+}.items():
+    QUANTMETHODS.register(names)(method)
+QUANTMETHODS.register("none", platform="musa")("no_quant.NoQuantization")

lightllm/common/quantization/w8a8.py

@@ -4,7 +4,6 @@
 import torch.nn.functional as F
 from typing import Optional, List, Union, Tuple
 from .quantize_method import QuantizationMethod, WeightPack
-from .registry import QUANTMETHODS
 from lightllm.common.basemodel.triton_kernel.quantization.scaled_mm_per_token_kernel import fp8_scaled_mm_per_token
 from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import (
     per_token_group_quant_fp8,
@@ -58,7 +57,6 @@ def _create_weight(
         raise NotImplementedError("Not implemented")
 
 
-@QUANTMETHODS.register(["w8a8-vllm", "w8a8"], platform="cuda")
 class w8a8QuantizationMethod(BaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -120,7 +118,6 @@ def _create_weight(
         return mm_param, mm_param_list
 
 
-@QUANTMETHODS.register(["fp8w8a8-vllm", "fp8w8a8"], platform="cuda")
 class FP8w8a8QuantizationMethod(BaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -320,7 +317,6 @@ def _create_weight_scale(self, out_dim: int, device_id: int) -> torch.Tensor:
         return torch.empty((1,), dtype=torch.float32, device=f"cuda:{device_id}")
 
 
-@QUANTMETHODS.register("fp8w8a8-pt-vllm", platform="cuda")
 class FP8w8a8PerTensorVllmQuantizationMethod(FP8w8a8PerTensorQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -349,7 +345,6 @@ def method_name(self):
         return "fp8w8a8-pt-vllm"
 
 
-@QUANTMETHODS.register("fp8w8a8-pt-sgl", platform="cuda")
 class FP8w8a8PerTensorSglQuantizationMethod(FP8w8a8PerTensorQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -382,7 +377,6 @@ def method_name(self):
         return "fp8w8a8-pt-sgl"
 
 
-@QUANTMETHODS.register(["fp8w8a8-pt-triton", "fp8w8a8-pt"], platform="cuda")
 class FP8w8a8PerTensorTritonQuantizationMethod(FP8w8a8PerTensorQuantizationMethod):
     def apply(
         self,
@@ -474,7 +468,6 @@ def _create_weight(
         return mm_param, mm_param_list
 
 
-@QUANTMETHODS.register("fp8w8a8-b128-vllm", platform="cuda")
 class FP8w8a8B128VllmQuantizationMethod(FP8w8a8B128QuantizationMethod):
     def apply(
         self,
@@ -502,7 +495,6 @@ def method_name(self):
         return "fp8w8a8-b128-vllm"
 
 
-@QUANTMETHODS.register(["fp8w8a8-b128-triton", "fp8w8a8-b128"], platform="cuda")
 class FP8w8a8B128TritonQuantizationMethod(FP8w8a8B128QuantizationMethod):
     def apply(
         self,

lightllm/common/quantization/w8a8gx.py

@@ -1,7 +1,6 @@
 import torch
 from typing import Optional, List, Union, Tuple
 from .quantize_method import QuantizationMethod
-from .registry import QUANTMETHODS
 
 from .quantize_method import WeightPack
 
@@ -37,7 +36,6 @@ def _create_weight(
         raise NotImplementedError("Not implemented")
 
 
-@QUANTMETHODS.register(["fp8w8a8g128-triton", "fp8w8a8g128"], platform="cuda")
 class FP8w8a8g128QuantizationMethod(_BaseQuantizationMethod):
     def __init__(self):
         super().__init__()
@@ -137,7 +135,6 @@ def _create_weight(
         return mm_param, mm_param_list
 
 
-@QUANTMETHODS.register(["fp8w8a8g64-triton", "fp8w8a8g64"], platform="cuda")
 class FP8w8a8g64QuantizationMethod(FP8w8a8g128QuantizationMethod):
     def __init__(self):
         super().__init__()

lightllm/server/api_cli.py

@@ -1,4 +1,5 @@
 import argparse
+from lightllm.common.quantization.registry import QUANTMETHODS
 
 
 def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
@@ -714,6 +715,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
     parser.add_argument(
         "--quant_type",
         type=str,
+        choices=QUANTMETHODS.supported_methods(platform=None),
         default="none",
         help=(
             "Quantization methods (W = weight, A = activation):\n"
@@ -758,6 +760,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
     parser.add_argument(
         "--vit_quant_type",
         type=str,
+        choices=QUANTMETHODS.supported_methods(platform=None),
         default="none",
         help="""Quantization method for ViT: w8a8 | fp8w8a8""",
     )

lightllm/server/api_server.py

@@ -1,12 +1,18 @@
 import torch
 from .api_cli import add_cli_args
+from lightllm.common.quantization import load_quant_config
 from lightllm.server.core.objs.start_args_type import StartArgs
 from lightllm.utils.log_utils import init_logger, set_log_node_role
 
 logger = init_logger(__name__)
 
 
 def launch_server(args: StartArgs):
+    for quant_type, quant_cfg in (("quant_type", "quant_cfg"), ("vit_quant_type", "vit_quant_cfg")):
+        try:
+            load_quant_config(getattr(args, quant_type), getattr(args, quant_cfg), platform=args.hardware_platform)
+        except ValueError as e:
+            raise ValueError(f"--{quant_type}/--{quant_cfg}: {e}") from e
     if args.use_dp_split_mode_connect_pd_master and args.run_mode != "prefill":
         raise ValueError("--use_dp_split_mode_connect_pd_master only supports prefill")
     set_log_node_role(args.run_mode)