from __future__ import annotations import logging from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, cast import torch from sglang.multimodal_gen.runtime.layers.linear import ( LinearMethodBase, UnquantizedLinearMethod, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, ) from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimW4A4Int4, ModelSlimW8A8Int8, ) if TYPE_CHECKING: from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimLinearScheme, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping logger = logging.getLogger(__name__) class ModelSlimConfig(QuantizationConfig): """ Config class for ModelSlim Quantization of Diffusion models https://gitcode.com/Ascend/msmodelslim, a NPU-specific quantization type. The quantization method (W8A8, W4A4, etc.) will be automatically parsed from the `quant_model_description.json` config. ModelSlim for Diffusion models includes support for various quantization schemes, such as: - W4A4 dynamic linear - W8A8 static linear - W8A8 dynamic linear """ def __init__( self, quant_config: Dict[str, Any] = {}, reverse_param_names_mapping: dict = None, ): super().__init__() self.quant_description = quant_config ignore = cast(List[str], quant_config.get("ignore", [])) self.ignore = ignore packed_modules_mapping = quant_config.get("packed_modules_mapping", {}) self.packed_modules_mapping = ( packed_modules_mapping if packed_modules_mapping is not None else {} ) self._name_mapper = ( get_param_names_mapping(reverse_param_names_mapping) if reverse_param_names_mapping is not None else None ) def get_linear_method(self) -> ModelSlimLinearMethod: return ModelSlimLinearMethod(self) @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: return [torch.int8, torch.float16, torch.bfloat16] @classmethod def get_min_capability(cls) -> int: return 0 @classmethod def get_name(cls) -> str: return "modelslim" @classmethod def get_config_filenames(cls) -> List[str]: filenames = ["quant_model_description.json"] return filenames @classmethod def from_config( cls, config: Dict[str, Any], reverse_param_names_mapping: dict = None ) -> ModelSlimConfig: return cls(config, reverse_param_names_mapping) def get_quant_method( self, layer: torch.nn.Module, prefix: str, ) -> Optional[QuantizeMethodBase]: from sglang.multimodal_gen.runtime.layers.linear import LinearBase if isinstance(layer, LinearBase): if should_ignore_layer( prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping, ): return UnquantizedLinearMethod() key = "model" packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {}) prefix_in_quant_config = prefix proj_name = prefix.split(".")[-1] if proj_name in packed_modules_mapping_subset: prefix_in_quant_config = prefix.replace( proj_name, packed_modules_mapping_subset[proj_name][0] ) if self.is_layer_skipped(prefix, packed_modules_mapping_subset): return UnquantizedLinearMethod() scheme = self.get_scheme(layer=layer, layer_name=prefix_in_quant_config) layer.scheme = scheme return ModelSlimLinearMethod(self) else: return None def _get_scheme_from_parts( self, layer_name: str, ) -> ModelSlimLinearScheme: full_weight_name = layer_name + ".weight" if self._name_mapper is not None: mapped_name, _, _ = self._name_mapper(full_weight_name) else: mapped_name = full_weight_name quant_type = self.quant_description.get(mapped_name, "") prefix = mapped_name.removesuffix(".weight") if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8": return ModelSlimW8A8Int8(quant_config=self.quant_description, prefix=prefix) elif quant_type == "W4A4_DYNAMIC": return ModelSlimW4A4Int4(quant_config=self.quant_description, prefix=prefix) elif quant_type == "W8A8_MXFP8": from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp8_scheme import ( ModelSlimMXFP8Scheme, ) return ModelSlimMXFP8Scheme() elif quant_type in ("W4A4_MXFP4", "W4A4_MXFP4_DUALSCALE"): from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp4_scheme import ( ModelSlimMXFP4Scheme, ) return ModelSlimMXFP4Scheme() raise NotImplementedError( f"No modelslim compatible scheme was found for layer '{layer_name}'. " f"quant_description['{layer_name}.weight'] = '{quant_type}'" ) def get_scheme( self, layer: torch.nn.Module, layer_name: Optional[str] = None ) -> Optional[ModelSlimLinearScheme]: """ get_scheme method adjusted for modelslim, taken from python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py """ scheme = self._get_scheme_from_parts( layer_name=layer_name, ) # Ascend doesn't support device capability logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) return scheme def is_layer_skipped( self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) ): # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped proj_name = prefix.split(".")[-1] if proj_name in fused_mapping: shard_prefixes = [ prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name] ] is_skipped = None for shard_prefix in shard_prefixes: is_shard_skipped = ( self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT" ) if is_skipped is None: is_skipped = is_shard_skipped elif is_shard_skipped != is_skipped: raise ValueError( f"Detected some but not all shards of {prefix} " "are quantized. All shards of fused layers " "to have the same precision." ) else: is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT" assert is_skipped is not None return is_skipped def get_scaled_act_names(self) -> List[str]: return [] class ModelSlimLinearMethod(LinearMethodBase): def __init__(self, quantization_config: ModelSlimConfig): self.quantization_config = quantization_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.scheme.process_weights_after_loading(layer) def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): """ Use the ModelSlimLinearScheme associated with each layer to create the necessary parameters for the layer. See LinearMethodBase for param details """ weight_loader = extra_weight_attrs.get("weight_loader") layer.scheme.create_weights( layer=layer, input_size=input_size, input_size_per_partition=input_size_per_partition, output_partition_sizes=output_partition_sizes, output_size=output_size, params_dtype=params_dtype, weight_loader=weight_loader, ) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ): """ Use the output of create_weights and the CompressedTensorsScheme associated with the layer to apply the forward pass with the layer input. See LinearMethodBase for param details """ scheme = layer.scheme if scheme is None: raise ValueError("A scheme must be defined for each layer") return scheme.apply_weights(layer, x, bias=bias)