add apply, serialize, deserialize support

Author: dsikka

src/compressed_tensors/base.py

@@ -18,3 +18,4 @@
 KV_CACHE_SCHEME_NAME = "kv_cache_scheme"
 COMPRESSION_VERSION_NAME = "version"
 QUANTIZATION_METHOD_NAME = "quant_method"
+TRANSFORMS_CONFIG = "transforms_config"

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -29,6 +29,7 @@
     QUANTIZATION_CONFIG_NAME,
     QUANTIZATION_METHOD_NAME,
     SPARSITY_CONFIG_NAME,
+    TRANSFORMS_CONFIG,
 )
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
@@ -38,13 +39,15 @@
     QuantizationStatus,
     apply_quantization_config,
     load_pretrained_quantization,
+    load_transforms,
 )
 from compressed_tensors.quantization.lifecycle import expand_target_names
 from compressed_tensors.quantization.quant_args import QuantizationArgs
 from compressed_tensors.quantization.utils import (
     is_module_quantized,
     iter_named_leaf_modules,
 )
+from compressed_tensors.transforms.transform_config import TransformationConfig
 from compressed_tensors.utils import (
     get_safetensors_folder,
     merge_names,
@@ -133,6 +136,8 @@ def from_compression_config(
 
         sparsity_config = cls.parse_sparsity_config(compression_config)
         quantization_config = cls.parse_quantization_config(compression_config)
+        transforms_config = cls.parse_transforms_config(compression_config)
+
         if sparsity_config is None and quantization_config is None:
             return None
 
@@ -144,8 +149,13 @@ def from_compression_config(
         if quantization_config is not None:
             quantization_config = QuantizationConfig.model_validate(quantization_config)
 
+        if transforms_config is not None:
+            transforms_config = TransformationConfig.model_validate(transforms_config)
+
         return cls(
-            sparsity_config=sparsity_config, quantization_config=quantization_config
+            sparsity_config=sparsity_config,
+            quantization_config=quantization_config,
+            transforms_config=transforms_config,
         )
 
     @classmethod
@@ -170,6 +180,10 @@ def from_pretrained_model(
             model, format=quantization_format
         )
 
+        # TODO: update to fetch from the pretrained model
+        # using the attached config for now
+        transforms_config = getattr(model, "transforms_config", None)
+
         if isinstance(sparsity_config, str):  # we passed in a sparsity format
             sparsity_config = SparsityCompressionConfig.load_from_registry(
                 sparsity_config
@@ -179,9 +193,25 @@ def from_pretrained_model(
             return None
 
         return cls(
-            sparsity_config=sparsity_config, quantization_config=quantization_config
+            sparsity_config=sparsity_config,
+            quantization_config=quantization_config,
+            transforms_config=transforms_config,
         )
 
+    @staticmethod
+    def parse_transforms_config(
+        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
+    ) -> Union[Dict[str, Any], None]:
+
+        if compression_config is None:
+            return None
+
+        if is_compressed_tensors_config(compression_config):
+            t_config = compression_config.transforms_config
+            return t_config.model_dump() if t_config is not None else None
+
+        return compression_config.get(TRANSFORMS_CONFIG, None)
+
     @staticmethod
     def parse_sparsity_config(
         compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
@@ -243,9 +273,11 @@ def __init__(
         self,
         sparsity_config: Optional[SparsityCompressionConfig] = None,
         quantization_config: Optional[QuantizationConfig] = None,
+        transforms_config: Optional[TransformationConfig] = None,
     ):
         self.sparsity_config = sparsity_config
         self.quantization_config = quantization_config
+        self.transforms_config = transforms_config
         self.sparsity_compressor = None
         self.quantization_compressor = None
 
@@ -434,10 +466,14 @@ def decompress(self, model_path: str, model: Module):
                 self.quantization_config, QuantizationStatus.FROZEN
             ):
                 names_to_scheme = apply_quantization_config(
-                    model, self.quantization_config
+                    model,
+                    self.quantization_config,
+                    transforms_config=self.transforms_config,
                 )
                 load_pretrained_quantization(model, model_path)
 
+            load_transforms(model, model_path)
+
             model_path_or_state_dict = (
                 model.state_dict() if sparse_decompressed else model_path
             )
@@ -497,6 +533,12 @@ def update_config(self, save_directory: str):
                 SPARSITY_CONFIG_NAME
             ] = sparsity_config_data
 
+        if self.transforms_config is not None:
+            transforms_config_data = self.transforms_config.to_dict()
+            config_data[QUANTIZATION_CONFIG_NAME][
+                TRANSFORMS_CONFIG
+            ] = transforms_config_data
+
         with open(config_file_path, "w") as config_file:
             json.dump(config_data, config_file, indent=2, sort_keys=True)
 

src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -126,6 +126,7 @@ def decompress_weight(
         :param quantization_args: quantization parameters for the weight
         :return: tensor of the decompressed weight
         """
+
         weight = compressed_data["weight_packed"]
         scale = compressed_data["weight_scale"]
         zero_point = compressed_data.get("weight_zero_point", None)

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -41,6 +41,9 @@
     iter_named_leaf_modules,
     iter_named_quantizable_modules,
 )
+from compressed_tensors.transforms import Transforms
+from compressed_tensors.transforms.transform_config import TransformationConfig
+from compressed_tensors.transforms.transform_data import TransformData
 from compressed_tensors.utils.helpers import fix_fsdp_module_name, replace_module
 from compressed_tensors.utils.offload import update_parameter_data
 from compressed_tensors.utils.safetensors_load import get_safetensors_folder
@@ -49,20 +52,45 @@
 
 __all__ = [
     "load_pretrained_quantization",
+    "load_transforms",
     "apply_quantization_config",
     "apply_quantization_status",
     "find_name_or_class_matches",
     "expand_target_names",
     "is_target",
+    "process_transforms_config",
 ]
 
 from compressed_tensors.quantization.utils.helpers import is_module_quantized
-from compressed_tensors.utils.safetensors_load import get_quantization_state_dict
+from compressed_tensors.utils.safetensors_load import (
+    get_quantization_state_dict,
+    get_weight_mappings,
+)
+from safetensors import safe_open
 
 
 _LOGGER = logging.getLogger(__name__)
 
 
+def load_transforms(model: Module, model_name_or_path: str):
+    model_path = get_safetensors_folder(model_name_or_path)
+    weight_mappings = get_weight_mappings(model_path)
+
+    state_dict = {}
+    for weight_name, safe_path in weight_mappings.items():
+        if "transform" in weight_name:
+            with safe_open(safe_path, framework="pt", device="cpu") as f:
+                state_dict[weight_name] = f.get_tensor(weight_name)
+
+    for name, submodule in iter_named_leaf_modules(model):
+        transform_data = getattr(submodule, "transform_data", None)
+        if transform_data:
+            for transform_name, transform_data in transform_data.data.items():
+                full_name = f"{name}.{transform_name}"
+                transform_data = state_dict.get(full_name, None)
+                update_parameter_data(submodule, transform_data, transform_name)
+
+
 def load_pretrained_quantization(model: Module, model_name_or_path: str):
     """
     Loads the quantization parameters (scale and zero point) from model_name_or_path to
@@ -104,8 +132,92 @@ def load_pretrained_quantization(model: Module, model_name_or_path: str):
             )
 
 
+def process_transforms_config(
+    transforms_config: TransformationConfig,
+    model: torch.nn.Module,
+    quantization_status: Optional[QuantizationStatus] = QuantizationStatus.INITIALIZED,
+):
+    for _, group in transforms_config.transform_groups.items():
+        # Each group/scheme targets one type of transform
+        transform_type = group.transform_type
+        transform_creation_args = group.transform_creation_args
+
+        # Need a better name - too many groups
+        for transform_arg in group.groups:
+            module_targets = transform_arg.module_targets
+
+            for name, submodule in model.named_modules():
+                if len(transform_arg.ignore) > 0:
+                    if matches := find_name_or_class_matches(
+                        name, submodule, transform_arg.ignore
+                    ):
+                        for match in matches:
+                            print("ignoring", match, name)
+                        continue  # layer matches ignore list, continue
+
+                targets = find_name_or_class_matches(
+                    name, submodule, transform_arg.targets
+                )
+
+                if targets:
+                    # Every layer which matches gets its own transform
+                    # Same transform type and args are used however
+
+                    # attach the transform to the submodule
+                    # because we can have more than one transform, need to attach some
+                    # form of key to fetch
+                    # OR we store it in the dictionary, handle cpu-offloading separatly
+
+                    if hasattr(submodule, "transform_data"):
+                        idx = submodule.transform_data.idx + 1
+                    else:
+                        idx = 0
+                    # only support weight parameters for now, assume one value in
+                    # module targets
+                    transform_name = f"{module_targets[0]}_transform_{idx}"
+
+                    # create an empty tensor OR create a new transform
+                    dtype = getattr(submodule, module_targets[0]).dtype
+                    if quantization_status in [
+                        QuantizationStatus.COMPRESSED,
+                        QuantizationStatus.FROZEN,
+                    ]:
+                        transform = Transforms.load_from_registry(
+                            transform_type,
+                            dtype=dtype,
+                            empty=True,
+                            **transform_creation_args,
+                        )
+                    else:
+                        transform = Transforms.load_from_registry(
+                            transform_type,
+                            dtype=dtype,
+                            **transform_creation_args,
+                        )
+                    setattr(submodule, transform_name, transform)
+
+                    # add relevant transform data to the submodule as well
+                    data = {
+                        transform_name: {
+                            "type": transform_type,
+                            "call_args": transform_arg.call_args,
+                        }
+                    }
+
+                    if hasattr(submodule, "transform_data"):
+                        submodule.transform_data.data.update(data)
+                        submodule.transform_data.idx = idx
+                    else:
+                        transform_data = TransformData(data=OrderedDict(data))
+                        submodule.transform_data = transform_data
+    return model
+
+
 def apply_quantization_config(
-    model: Module, config: Union[QuantizationConfig, None], run_compressed: bool = False
+    model: Module,
+    config: Union[QuantizationConfig, None],
+    run_compressed: bool = False,
+    transforms_config=None,
 ) -> OrderedDict:
     """
     Initializes the model for quantization in-place based on the given config.
@@ -184,6 +296,12 @@ def apply_quantization_config(
                 f"{set(config.ignore) - set(ignored_submodules)}"
             )
 
+    if transforms_config:
+        model.transforms_config = transforms_config
+        model = process_transforms_config(
+            transforms_config, model, config.quantization_status
+        )
+
     # apply current quantization status across all targeted layers
     apply_quantization_status(model, config.quantization_status)
     return names_to_scheme

src/compressed_tensors/quantization/lifecycle/forward.py

@@ -28,6 +28,10 @@
     calculate_range,
     compute_dynamic_scales_and_zp,
 )
+from compressed_tensors.transforms.apply import (
+    apply_inverse_transforms_to_parameter,
+    apply_transforms_to_parameter,
+)
 from compressed_tensors.utils import safe_permute
 from torch.nn import Module
 
@@ -280,10 +284,25 @@ def wrapped_forward(self, *args, **kwargs):
         if scheme.weights is not None and not compressed:
             # calibrate and (fake) quantize weights when applicable
             unquantized_weight = self.weight.data.clone()
+            transform_data = getattr(module, "transform_data", None)
+            if transform_data is not None:
+                apply_transforms_to_parameter(
+                    module=module,
+                    module_parameter=self.weight,
+                    transform_data=transform_data,
+                )
+
             self.weight.data = forward_quantize(
                 module, self.weight, "weight", scheme.weights
             )
 
+            if transform_data is not None:
+                apply_inverse_transforms_to_parameter(
+                    module=module,
+                    module_parameter=self.weight,
+                    transform_data=transform_data,
+                )
+
         # perform wrapped forward call
         output = forward_func_orig.__get__(module, module.__class__)(
             input_, *args[1:], **kwargs