chore: import upstream snapshot with attribution

python/mlc_llm/quantization/model_quantization.py

@@ -0,0 +1,157 @@
+"""Quantization factory utilities for model quantization."""
+
+from typing import Any, Callable, Dict, Optional, Tuple, Type  # noqa: UP035
+
+from tvm.relax.frontend import nn
+
+from mlc_llm.loader import QuantizeMapping
+
+from .awq_quantization import AWQQuantize
+from .block_scale_quantization import BlockScaleQuantize
+from .ft_quantization import FTQuantize
+from .group_quantization import GroupQuantize
+from .no_quantization import NoQuantize
+from .per_tensor_quantization import PerTensorQuantize
+from .quantization import Quantization
+
+FuncQuantization = Callable[[Any, Quantization], Tuple[nn.Module, QuantizeMapping]]  # noqa: UP006
+
+
+def make_quantization_functions(
+    model_cls: Type[nn.Module],  # noqa: UP006
+    *,
+    model_ctor: Optional[Callable[[Any], nn.Module]] = None,
+    supports_group_quant: bool = True,
+    supports_ft_quant: bool = True,
+    supports_awq: bool = False,
+    awq_unsupported_message: Optional[str] = None,
+    supports_per_tensor: bool = False,
+    supports_block_scale: bool = False,
+    set_tensor_parallel_shards: bool = True,
+    per_tensor_use_shards: bool = True,
+) -> Dict[str, FuncQuantization]:  # noqa: UP006
+    """Create standard quantization function implementations for a model class."""
+
+    def _create_model(model_config: Any) -> nn.Module:
+        if model_ctor is not None:
+            return model_ctor(model_config)
+        return model_cls(model_config)
+
+    def _no_quant(model_config: Any, quantization: NoQuantize) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        return model, QuantizeMapping({}, {})
+
+    def _group_quant(
+        model_config: Any,
+        quantization: GroupQuantize,
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        if set_tensor_parallel_shards:
+            if not hasattr(model_config, "tensor_parallel_shards"):
+                raise AttributeError(
+                    "model_config is missing required "
+                    "attribute 'tensor_parallel_shards' for group quantization"
+                )
+            quantization.tensor_parallel_shards = getattr(model_config, "tensor_parallel_shards")
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    def _ft_quant(model_config: Any, quantization: FTQuantize) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    def _awq_quant(
+        model_config: Any, quantization: AWQQuantize
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        if awq_unsupported_message is not None:
+            raise NotImplementedError(awq_unsupported_message)
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    def _per_tensor_quant(
+        model_config: Any,
+        quantization: PerTensorQuantize,
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        kwargs = {}
+        if per_tensor_use_shards:
+            if not hasattr(model_config, "tensor_parallel_shards"):
+                raise AttributeError(
+                    "model_config is missing required attribute "
+                    "'tensor_parallel_shards' for per-tensor quantization"
+                )
+            kwargs["tensor_parallel_shards"] = getattr(model_config, "tensor_parallel_shards")
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+            **kwargs,
+        )
+        return model, quant_map
+
+    def _block_scale_quant(
+        model_config: Any,
+        quantization: BlockScaleQuantize,
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(model, quant_map, "")
+        return model, quant_map
+
+    quantize_fns: Dict[str, FuncQuantization] = {"no-quant": _no_quant}  # noqa: UP006
+    if supports_group_quant:
+        quantize_fns["group-quant"] = _group_quant
+    if supports_ft_quant:
+        quantize_fns["ft-quant"] = _ft_quant
+    if supports_awq:
+        quantize_fns["awq"] = _awq_quant
+    if supports_per_tensor:
+        quantize_fns["per-tensor-quant"] = _per_tensor_quant
+    if supports_block_scale:
+        quantize_fns["block-scale-quant"] = _block_scale_quant
+    return quantize_fns
+
+
+def make_awq_quant(
+    model_cls: Type[nn.Module],  # noqa: UP006
+) -> Callable[[Any, AWQQuantize], Tuple[nn.Module, QuantizeMapping]]:  # noqa: UP006
+    """Create a standard AWQ quantization function for loaders."""
+
+    def awq_quant(
+        model_config: Any, quantization: AWQQuantize
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = model_cls(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    return awq_quant