225 lines
7.7 KiB
Python
225 lines
7.7 KiB
Python
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.modelslim import ModelSlimConfig
|
|
from sglang.srt.layers.quantization.modelslim.schemes import (
|
|
ModelSlimLinearScheme,
|
|
)
|
|
|
|
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] = {}):
|
|
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 {}
|
|
)
|
|
|
|
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]) -> ModelSlimConfig:
|
|
return cls(config)
|
|
|
|
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:
|
|
|
|
quant_type = self.quant_description.get(layer_name + ".weight", "")
|
|
if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8":
|
|
return ModelSlimW8A8Int8(
|
|
quant_config=self.quant_description, prefix=layer_name
|
|
)
|
|
elif quant_type == "W4A4_DYNAMIC":
|
|
return ModelSlimW4A4Int4(
|
|
quant_config=self.quant_description, prefix=layer_name
|
|
)
|
|
raise NotImplementedError("No modelslim compatible scheme was found.")
|
|
|
|
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)
|