angelslim 0.1.0__py3-none-any.whl

angelslim/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .engine import Engine  # noqa: F401

angelslim/compressor/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .compressor_factory import CompressorFactory  # noqa: F401

angelslim/compressor/compressor_factory.py

@@ -0,0 +1,83 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Callable, Dict, Optional, Type, Union
+
+from ..utils import print_info
+
+
+class CompressorFactory:
+    """
+    Factory class for model compression methods with flexible registration.
+    Supports both explicit name registration and direct class name registration.
+    """
+
+    _compress_methods: Dict[str, Type[Any]] = {}
+
+    @classmethod
+    def register(cls, name: Optional[Union[str, Callable]] = None) -> Callable:
+        """Decorator to register compression methods. Supports two usage patterns:
+        1. @CompressorFactory.register("explicit_name")
+        2. @CompressorFactory.register (uses class name as key)
+        """
+
+        # Handler for direct class registration (@CompressorFactory.register)
+        def register_class(compress_cls: Type[Any]) -> Type[Any]:
+            """Register a class using its own name as the key"""
+            key = compress_cls.__name__
+            if key in cls._compress_methods:
+                print_info(
+                    f"Compression method '{key}' already exists, will be overwritten."
+                )
+            cls._compress_methods[key] = compress_cls
+            return compress_cls
+
+        # Handler for named registration (@CompressorFactory.register("name"))
+        def register_with_name(key: str) -> Callable[[Type[Any]], Type[Any]]:
+            """Decorator that registers a class with a custom key"""
+
+            def decorator(compress_cls: Type[Any]) -> Type[Any]:
+                if key in cls._compress_methods:
+                    print_info(f"register '{key}' already exists, will be overwritten.")
+                cls._compress_methods[key] = compress_cls
+                return compress_cls
+
+            return decorator
+
+        # Determine registration type based on input
+        if name is None:
+            # Case 1: Direct class registration (@CompressorFactory.register)
+            return register_class
+        elif isinstance(name, str):
+            # Case 2: Explicit name registration (@CompressorFactory.register("name"))
+            return register_with_name(name)
+        elif callable(name):
+            # Case 3: Direct class registration (called without parentheses)
+            return register_class(name)
+        else:
+            raise TypeError("Invalid argument type for registration")
+
+    @classmethod
+    def create(cls, name: str, model: Any, slim_config: Any) -> Any:
+        """Create compressor instance"""
+        if name not in cls._compress_methods:
+            available = list(cls._compress_methods.keys())
+            raise ValueError(
+                f"Compression method '{name}' not registered. Available: {available}"
+            )
+        return cls._compress_methods[name](model, slim_config)
+
+    @classmethod
+    def get_available_compressor(cls) -> list:
+        return list(cls._compress_methods.keys())

angelslim/compressor/distill/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

angelslim/compressor/quant/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .core import *  # noqa: F401 F403
+from .modules import *  # noqa: F401 F403
+from .observers import *  # noqa: F401 F403
+from .ptq import PTQ  # noqa: F401

angelslim/compressor/quant/core/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .config import *  # noqa: F401 F403
+from .hook import DiTHook, PTQHook  # noqa: F401
+from .metrics import mse_loss, snr_loss  # noqa: F401
+from .packing_utils import dequantize_gemm  # noqa: F401
+from .quant_func import *  # noqa: F401 F403
+from .sample_func import EMASampler, MultiStepSampler  # noqa: F401
+from .save import PTQPTMSave  # noqa: F401
+from .save import PTQSaveVllmHF  # noqa: F401
+from .save import PTQTorchSave  # noqa: F401
+from .save import PTQvLLMSaveHF  # noqa: F401
+from .save import PTQVLMSaveVllmHF  # noqa: F401

angelslim/compressor/quant/core/config.py

@@ -0,0 +1,166 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from ..observers import (
+    AbsMaxChannelWiseWeightObserver,
+    AbsMaxGroupWiseWeightObserver,
+    AbsmaxPerchannelObserver,
+    AbsmaxPertensorObserver,
+)
+
+ACT_OBSERVERS_CLASS = {
+    "per-tensor": AbsmaxPertensorObserver,
+    "per-channel": AbsmaxPerchannelObserver,
+}
+WEIGHT_OBSERVERS_CLASS = {
+    "per-tensor": AbsmaxPertensorObserver,
+    "per-channel": AbsMaxChannelWiseWeightObserver,
+    "per-group": AbsMaxGroupWiseWeightObserver,
+}
+
+KVCACHE_OBSERVERS_CLASS = {"per-channel": AbsmaxPerchannelObserver}
+
+
+class QuantConfig:
+    r"""
+    Configure how to quantize a model or a part of the model. It will map each layer to
+    an instance of observers by the settings.
+
+    Args:
+        config: The quant config.
+    Examples:
+        .. code-block:: python
+
+            >>> from slim.quant import QuantConfig
+            >>> q_config = QuantConfig(yaml_config)
+    """
+
+    def __init__(self, config, global_config=None):
+        # quant_algo change
+        self.act_observer = None
+        self.weight_observer = None
+        self.kv_cache_observer = None
+
+        quantization_args = config.quantization
+        self.quant_algo = quantization_args.name
+        self.quant_bit = quantization_args.bits
+        self.max_seq_length = global_config.max_seq_length
+        self.quant_helpers = quantization_args.quant_helpers
+        act_quant_method = quantization_args.quant_method.get("activation", None)
+        weight_quant_method = quantization_args.quant_method["weight"]
+
+        if "fp8" in self.quant_algo:
+            is_dynamic = "dynamic" if "dynamic" in self.quant_algo else "static"
+            assert (
+                is_dynamic or act_quant_method is not None
+            ), "[OpenSlim][Error] fp8_static need act_quant_method"
+            self.act_observer = (
+                ACT_OBSERVERS_CLASS[act_quant_method]
+                if "static" in is_dynamic
+                else None
+            )
+            self.weight_observer = WEIGHT_OBSERVERS_CLASS[weight_quant_method]
+            self.kv_cache_observer = None
+            self.quant_algo_info = {
+                "w": f"fp8_{weight_quant_method}",
+                "ignore_layers": quantization_args.ignore_layers,
+            }
+            if act_quant_method is not None:
+                self.quant_algo_info["a"] = f"fp8_{act_quant_method}-{is_dynamic}"
+            self.hidden_size = global_config.hidden_size
+            self.model_arch_type = global_config.model_arch_type
+            self.low_memory = config.quantization.low_memory
+        elif "int8" in self.quant_algo:
+            is_dynamic = "dynamic" if "dynamic" in self.quant_algo else "static"
+            assert (
+                is_dynamic or act_quant_method is not None
+            ), "[OpenSlim][Error] int8_static need act_quant_method"
+            self.act_observer = (
+                ACT_OBSERVERS_CLASS[act_quant_method]
+                if "static" in is_dynamic
+                else None
+            )
+            self.weight_observer = WEIGHT_OBSERVERS_CLASS[weight_quant_method]
+            self.kv_cache_observer = None
+            self.quant_algo_info = {
+                "w": f"int8_{weight_quant_method}",
+                "ignore_layers": quantization_args.ignore_layers,
+            }
+            if act_quant_method is not None:
+                self.quant_algo_info["a"] = f"int8_{act_quant_method}-{is_dynamic}"
+            self.hidden_size = global_config.hidden_size
+            self.model_arch_type = global_config.model_arch_type
+            self.low_memory = config.quantization.low_memory
+        elif "int4_awq" in self.quant_algo:
+            self.act_observer = None
+            self.weight_observer = None
+            self.kv_cache_observer = None
+            group_size = (
+                128
+                if quantization_args.quant_method["group_size"] == -1
+                else quantization_args.quant_method["group_size"]
+            )
+            self.quant_algo_info = {
+                "zero_point": quantization_args.quant_method["zero_point"],
+                "group_size": int(group_size),
+                "mse_range": quantization_args.quant_method["mse_range"],
+            }
+            self.hidden_size = global_config.hidden_size
+            self.model_arch_type = global_config.model_arch_type
+        elif "int4_gptq" in self.quant_algo:
+            self.act_observer = None
+            self.weight_observer = None
+            self.kv_cache_observer = None
+            group_size = (
+                128
+                if quantization_args.quant_method["group_size"] == -1
+                else quantization_args.quant_method["group_size"]
+            )
+            self.quant_algo_info = {
+                "group_size": group_size,
+                "ignore_layers": quantization_args.ignore_layers,
+            }
+            self.hidden_size = global_config.hidden_size
+
+        if "smooth" in self.quant_helpers:
+            self.smooth_alpha = quantization_args.smooth_alpha
+            self.smooth_observer = ACT_OBSERVERS_CLASS["per-channel"]
+        self.custom_observe_layers_names = "default"
+
+    def custom_observe_layers(
+        self,
+        names: List,
+        act_observer="default",
+        weight_observer="default",
+        kv_cache_observer="default",
+    ):
+        """
+        name supports fuzzy search.
+        """
+        self.custom_observe_layers_names = names
+        self.act_observer = (
+            act_observer if act_observer in ACT_OBSERVERS_CLASS else self.act_observer
+        )
+        self.weight_observer = (
+            weight_observer
+            if weight_observer in WEIGHT_OBSERVERS_CLASS
+            else self.weight_observer
+        )
+        self.kv_cache_observer = (
+            kv_cache_observer
+            if kv_cache_observer in KVCACHE_OBSERVERS_CLASS
+            else self.kv_cache_observer
+        )

angelslim/compressor/quant/core/dit_hook.py

@@ -0,0 +1,92 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+import torch
+import torch.nn as nn
+
+
+def filter_func(name):
+    pattern = re.compile(
+        r".*(mlp_t5|pooler|style_embedder|x_embedder|t_embedder|extra_embedder).*"
+    )
+    return pattern.match(name) is not None
+
+
+class DiTHook:
+    def __init__(self, model, use_transformer_engine=False):
+        """
+        Args:
+            model(nn.Moudle, required): the model to be quant
+        """
+        self.model = model
+        self.input_activation = {}
+        self.output_activation = {}
+        self.input_activation_cnt = {}
+        self.output_activation_cnt = {}
+        self.use_transformer_engine = use_transformer_engine
+        self._apply_hook()
+
+    def _apply_hook(self):
+        self._forward_hook_list = []
+        for name, sub_layer in self.model.named_modules():
+            if filter_func(name):
+                continue
+            instance_list = nn.Linear
+            if isinstance(sub_layer, instance_list):
+                forward_pre_hook_handle = sub_layer.register_forward_hook(
+                    self._forward_pre_hook
+                )
+                self._forward_hook_list.append(forward_pre_hook_handle)
+
+    def _forward_pre_hook(self, layer, input, output):
+        layer_name = ""
+        for name, module in self.model.named_modules():
+            if filter_func(name):
+                continue
+            if module == layer:
+                layer_name = name
+                break
+        x = (
+            output[0].detach().cpu()
+            if isinstance(output, tuple)
+            else output.detach().cpu()
+        )
+        self.output_activation[layer_name] = (
+            self.output_activation.get(layer_name, torch.zeros(x.shape).to(x.dtype)) + x
+        )
+        self.output_activation_cnt[layer_name] = (
+            self.output_activation_cnt.get(layer_name, 0) + 1
+        )
+        y = (
+            input[0].detach().cpu()
+            if isinstance(input, tuple)
+            else input.detach().cpu()
+        )
+        self.input_activation[layer_name] = (
+            self.input_activation.get(layer_name, torch.zeros(y.shape).to(y.dtype)) + y
+        )
+        self.input_activation_cnt[layer_name] = (
+            self.input_activation_cnt.get(layer_name, 0) + 1
+        )
+
+    def remove_hook(self):
+        for hook in self._forward_hook_list:
+            hook.remove()
+        self._forward_hook_list = []
+
+    def clean_acitvation_list(self):
+        self.input_activation = {}
+        self.output_activation = {}

angelslim/compressor/quant/core/hook.py

@@ -0,0 +1,184 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+import torch
+
+from ..observers import ParentObserver, PTQObserver
+from .quant_func import get_fp_maxval, get_fp_search_maxval
+
+__all__ = ["PTQHook", "DiTHook"]
+
+
+class PTQHook:
+    def __init__(self, model):
+        self.quant_model = model
+        self._forward_hook_list = []
+        # {name: layer}
+        self.quant_layers_dict = {}
+        # {layer: observer}
+        self.observer_dict = {}
+        self.kv_names = []
+
+    def apply_hook(self):
+        self.quant_layers_dict = self.quant_model.get_observer_layers()
+        self.kv_names = self.quant_model.get_kvcache_observer_layers_names(
+            self.quant_layers_dict.keys()
+        )
+        act_observer = self.quant_model.quant_algo_dict["act_observer"]
+        weight_observer = self.quant_model.quant_algo_dict["weight_observer"]
+        kv_cache_observer = self.quant_model.quant_algo_dict["kv_cache_observer"]
+
+        quant_parent_dict = self.quant_model.get_parent_dict(self.quant_layers_dict)
+        parent_observers = {
+            v: ParentObserver() for v in set(quant_parent_dict.values())
+        }
+
+        # apply observers
+        for name, sub_layer in self.quant_layers_dict.items():
+            extra_kwargs = (
+                {"parent_observer": parent_observers[quant_parent_dict[name]]}
+                if name in quant_parent_dict
+                else {}
+            )
+            observer = PTQObserver(
+                sub_layer,
+                act_observer,
+                weight_observer,
+                kv_cache_observer if name in self.kv_names else None,
+                self.quant_model.quant_algo_dict,
+                **extra_kwargs
+            )
+            forward_hook_handle = sub_layer.register_forward_hook(self._forward_hook)
+            self.observer_dict[sub_layer] = observer
+            self._forward_hook_list.append(forward_hook_handle)
+
+    def apply_smooth_hook(self, smooth_mapping_layers, smooth_observer):
+        for smooth_layer, _ in smooth_mapping_layers.values():
+            observer = PTQObserver(
+                smooth_layer,
+                act_observer=None,
+                weight_observer=None,
+                kv_cache_observer=None,
+                quant_algo_dict=self.quant_model.quant_algo_dict,
+                smooth_act_observer=smooth_observer,
+            )
+            forward_hook_handle = smooth_layer.register_forward_hook(self._forward_hook)
+            self.observer_dict[smooth_layer] = observer
+            self._forward_hook_list.append(forward_hook_handle)
+
+    def _forward_hook(self, layer, input, output):
+        x = input[0].clone() if isinstance(input, tuple) else input.clone()
+        y = output[0].clone() if isinstance(output, tuple) else output.clone()
+        if hasattr(self.quant_model, "apply_layer_norm_list"):
+            if layer in self.quant_model.apply_layer_norm_list:
+                x = self.quant_model.apply_layer_norm(layer, x)
+        self.observer_dict[layer](x, y)
+        return output
+
+    def remove_hook(self):
+        for hook in self._forward_hook_list:
+            hook.remove()
+        self._forward_hook_list = []
+
+    def post_process(self):
+        maxval = get_fp_maxval(bits=8)
+        if self.quant_model.quant_algo_dict["w_quant_algo"] == "fp8":
+            for k, v in self.quant_model.weight_scales_dict.items():
+                self.quant_model.weight_scales_dict[k] = v / maxval.type(v.dtype)
+        if self.quant_model.quant_algo_dict["a_quant_algo"] == "fp8":
+            for name, sub_layer in self.quant_layers_dict.items():
+                if sub_layer in self.observer_dict:
+                    if name in self.quant_model.act_scales_dict.keys():
+                        act_dtype = self.quant_model.act_scales_dict[name].dtype
+                        if "Search" in str(self.observer_dict[sub_layer]):
+                            tmp_maxval = get_fp_search_maxval(
+                                self.observer_dict[sub_layer].sampled_input
+                            )
+                            self.quant_model.act_scales_dict[name] = (
+                                self.quant_model.act_scales_dict[name]
+                                / tmp_maxval.type(act_dtype)
+                            )
+                        else:
+                            self.quant_model.act_scales_dict[name] = (
+                                self.quant_model.act_scales_dict[name]
+                                / maxval.type(act_dtype)
+                            )
+        if self.quant_model.quant_algo_dict["c_quant_algo"] == "fp8":
+            for k, v in self.quant_model.kv_cache_scales_dict.items():
+                self.quant_model.kv_cache_scales_dict[k] = v / maxval.type(v.dtype)
+
+
+def _filter_func(name):
+    pattern = re.compile(
+        r".*(mlp_t5|pooler|style_embedder|x_embedder|t_embedder|extra_embedder).*"
+    )
+    return pattern.match(name) is not None
+
+
+class DiTHook:
+    def __init__(self, model):
+        """
+        Args:
+            model(nn.Moudle, required): the model to be quant
+        """
+        self.model = model
+        self.input_activation = []
+        self.output_activation = []
+
+        self._apply_hook()
+
+    def _apply_hook(self):
+        self._forward_hook_list = []
+        for name, sub_layer in self.model.named_modules():
+            if _filter_func(name):
+                continue
+            if isinstance(sub_layer, (torch.nn.Conv2d, torch.nn.Linear)):
+                if "blocks" in name:
+                    # handle
+                    forward_pre_hook_handle = sub_layer.register_forward_hook(
+                        self._forward_pre_hook
+                    )
+                    self._forward_hook_list.append(forward_pre_hook_handle)
+
+    def _forward_pre_hook(self, layer, input, output):
+        layer_name = ""
+        for name, module in self.model.named_modules():
+            if _filter_func(name):
+                continue
+            if module == layer:
+                layer_name = name
+                break
+        x = (
+            output[0].detach().cpu()
+            if isinstance(output, tuple)
+            else output.detach().cpu()
+        )
+        self.output_activation.append((layer_name, x))
+        y = (
+            input[0].detach().cpu()
+            if isinstance(input, tuple)
+            else input.detach().cpu()
+        )
+        self.input_activation.append((layer_name, y))
+
+    def remove_hook(self):
+        for hook in self._forward_hook_list:
+            hook.remove()
+        self._forward_hook_list = []
+
+    def clean_acitvation_list(self):
+        self.input_activation = []
+        self.output_activation = []