Merge branch 'kylesayrs/serialize-tconfig' into kylesayrs/transform-merge

Author: kylesayrs

src/compressed_tensors/base.py

@@ -12,9 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-SPARSITY_CONFIG_NAME = "sparsity_config"
+# configs
 QUANTIZATION_CONFIG_NAME = "quantization_config"
-COMPRESSION_CONFIG_NAME = "compression_config"
-KV_CACHE_SCHEME_NAME = "kv_cache_scheme"
+SPARSITY_CONFIG_NAME = "sparsity_config"
+TRANSFORM_CONFIG_NAME = "transform_config"
+
+# required fields
 COMPRESSION_VERSION_NAME = "version"
 QUANTIZATION_METHOD_NAME = "quant_method"
+
+# auxillary configs
+KV_CACHE_SCHEME_NAME = "kv_cache_scheme"

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -29,6 +29,7 @@
     QUANTIZATION_CONFIG_NAME,
     QUANTIZATION_METHOD_NAME,
     SPARSITY_CONFIG_NAME,
+    TRANSFORM_CONFIG_NAME,
 )
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.compressors.sparse_compressors import DenseCompressor
@@ -43,6 +44,7 @@
 )
 from compressed_tensors.quantization.lifecycle import expand_target_names
 from compressed_tensors.quantization.utils import is_module_quantized
+from compressed_tensors.transform import TransformConfig
 from compressed_tensors.utils import (
     align_module_device,
     delete_offload_parameter,
@@ -105,6 +107,7 @@ class ModelCompressor:
 
     sparsity_config: Optional[SparsityCompressionConfig] = None
     quantization_config: Optional[QuantizationConfig] = None
+    transform_config: Optional[TransformConfig] = None
 
     @classmethod
     def from_pretrained(
@@ -144,6 +147,8 @@ def from_compression_config(
 
         sparsity_config = cls.parse_sparsity_config(compression_config)
         quantization_config = cls.parse_quantization_config(compression_config)
+        # NOTE: transfrom config is not support by ctconfig yet
+
         if sparsity_config is None and quantization_config is None:
             return None
 
@@ -177,20 +182,27 @@ def from_pretrained_model(
             algorithm
         :return: compressor for the configs, or None if model is not compressed
         """
+        # reconstruct config from schemes attached to modules
         quantization_config = QuantizationConfig.from_pretrained(
             model, format=quantization_format
         )
 
+        # use config passed as argument
         if isinstance(sparsity_config, str):  # we passed in a sparsity format
             sparsity_config = SparsityCompressionConfig.load_from_registry(
                 sparsity_config
             )
 
-        if sparsity_config is None and quantization_config is None:
+        # use config attached to model
+        transform_config = getattr(model, TRANSFORM_CONFIG_NAME, None)
+
+        if not any((quantization_config, sparsity_config, transform_config)):
             return None
 
         return cls(
-            sparsity_config=sparsity_config, quantization_config=quantization_config
+            sparsity_config=sparsity_config,
+            quantization_config=quantization_config,
+            transform_config=transform_config,
         )
 
     @staticmethod
@@ -254,13 +266,17 @@ def __init__(
         self,
         sparsity_config: Optional[SparsityCompressionConfig] = None,
         quantization_config: Optional[QuantizationConfig] = None,
+        transform_config: Optional[TransformConfig] = None,
     ):
         self.sparsity_config = sparsity_config
         self.quantization_config = quantization_config
+        self.transform_config = transform_config
+
         self.sparsity_compressor = None
         self.quantization_compressor: Optional[
             Union[BaseQuantizationCompressor, DenseCompressor]
         ] = None
+        # no transform compressor is required
 
         if sparsity_config is not None:
             self.sparsity_compressor = BaseCompressor.load_from_registry(
@@ -640,43 +656,49 @@ def update_config(self, save_directory: str):
 
         :param save_directory: path to a folder containing a HF model config
         """
-        if self.quantization_config is None and self.sparsity_config is None:
+        # this check is also done in `from_pretrained_model`,
+        # but not in `from_pretrained`` or `from_compression_config``
+        if not any(
+            (self.quantization_config, self.sparsity_config, self.transform_config)
+        ):
             return
 
+        # write to config.json file, regardless of whether it exists already
+        # overwrite previous config and version if already existing
         config_file_path = os.path.join(save_directory, CONFIG_NAME)
-        if not os.path.exists(config_file_path):
-            _LOGGER.warning(
-                f"Could not find a valid model config file in "
-                f"{save_directory}. Compression config will not be saved."
-            )
-            return
+        if os.path.exists(config_file_path):
+            with open(config_file_path, "r") as file:
+                config_data = json.load(file)
+        else:
+            config_data = {}
 
-        with open(config_file_path, "r") as config_file:
-            config_data = json.load(config_file)
+        # serialize configs into json
+        qconfig_data = (
+            self.quantization_config.model_dump(exclude=["quant_method", "format"])
+            if self.quantization_config is not None
+            else {}
+        )
+        sconfig_data = (
+            self.sparsity_config.model_dump()
+            if self.sparsity_config is not None
+            else {}
+        )
+        tconfig_data = (
+            self.transform_config.model_dump()
+            if self.transform_config is not None
+            else {}
+        )
 
-        # required metadata whenever a quantization or sparsity config is present
-        # overwrite previous config and version if already existing
-        config_data[QUANTIZATION_CONFIG_NAME] = {}
-        config_data[QUANTIZATION_CONFIG_NAME][
-            COMPRESSION_VERSION_NAME
-        ] = compressed_tensors.__version__
-        if self.quantization_config is not None:
-            self.quantization_config.quant_method = DEFAULT_QUANTIZATION_METHOD
-        else:
-            config_data[QUANTIZATION_CONFIG_NAME][
-                QUANTIZATION_METHOD_NAME
-            ] = DEFAULT_QUANTIZATION_METHOD
-
-        # quantization and sparsity configs
-        if self.quantization_config is not None:
-            quant_config_data = self.quantization_config.model_dump()
-            config_data[QUANTIZATION_CONFIG_NAME] = quant_config_data
-        if self.sparsity_config is not None:
-            sparsity_config_data = self.sparsity_config.model_dump()
-            config_data[QUANTIZATION_CONFIG_NAME][
-                SPARSITY_CONFIG_NAME
-            ] = sparsity_config_data
+        # construct compression (quantization) config
+        config_data[QUANTIZATION_CONFIG_NAME] = {
+            COMPRESSION_VERSION_NAME: compressed_tensors.__version__,
+            QUANTIZATION_METHOD_NAME: DEFAULT_QUANTIZATION_METHOD,
+            SPARSITY_CONFIG_NAME: sconfig_data,
+            TRANSFORM_CONFIG_NAME: tconfig_data,
+            **qconfig_data,
+        }
 
+        # write results to config.json file
         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/nvfp4_quantized.py

@@ -61,6 +61,27 @@ def compression_param_names(self) -> Tuple[str]:
             "weight_global_scale",
         )
 
+    def compression_param_info(
+        self,
+        weight_shape: torch.Size,
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        """
+        Creates a dictionary of expected shapes and dtypes for each compression
+            parameter used by the compressor
+
+        :param weight_shape: uncompressed weight shape
+        :param quantization_args: quantization parameters for the weight
+        :return: dictionary mapping compressed parameter names to shape and dtype
+        """
+        output = {
+            "weight_packed": (
+                torch.Size((weight_shape[0], weight_shape[1] // 2)),
+                torch.uint8,
+            ),
+        }
+        return output
+
     def compress_weight(
         self,
         weight: Tensor,

src/compressed_tensors/quantization/lifecycle/forward.py

@@ -257,13 +257,10 @@ def _process_quantization(
         QuantizationStrategy.GROUP,
         QuantizationStrategy.TENSOR_GROUP,
     ):
-        n_dims = x.shape
-        if len(n_dims) > 2:
-            x = x.squeeze(0)
 
         output_dtype = dtype if dtype is not None else x.dtype
         output = torch.zeros_like(x).to(output_dtype)
-        columns = output.shape[1]
+        columns = output.shape[-1]
 
         # TODO: make validation step for inputs
 
@@ -293,14 +290,12 @@ def _process_quantization(
             perm = torch.argsort(g_idx)
             x = safe_permute(x, perm, dim=1)
 
-        x = torch.reshape(
-            x,
-            (
-                x.shape[0],
-                ceil(x.shape[1] / group_size),
-                group_size,
-            ),
+        # Maintain all dimensions apart from the last dim, which is divided by the group_size
+        reshaped_dims = (
+            ceil(x.shape[-1] / group_size),
+            group_size,
         )
+        x = x.unflatten(-1, reshaped_dims)
 
         if do_quantize:
             output = _quantize(
@@ -323,19 +318,12 @@ def _process_quantization(
                 global_scale=global_scale,
             )
 
-        output = torch.reshape(
-            output,
-            (output.shape[0], output.shape[1] * output.shape[2]),
-        )
-
+        output = output.flatten(start_dim=-2)
         output = output.to(output_dtype)
 
         if not is_column_order:
             output = safe_permute(output, torch.argsort(perm), dim=1)
 
-        if len(n_dims) > 2:
-            output = output.unsqueeze(0)
-
     else:  # covers channel, token and tensor strategies
         if do_quantize:
             output = _quantize(

src/compressed_tensors/quantization/utils/helpers.py

@@ -175,20 +175,16 @@ def compute_dynamic_scales_and_zp(
         QuantizationStrategy.TENSOR_GROUP,
         QuantizationStrategy.GROUP,
     ):
-        if len(value.shape) > 2:
-            value = value.squeeze(0)
 
-        dim = {0, 1}
-        reduce_dims = tuple(idx for idx in range(3) if idx not in dim)
+        reduce_dims = -1
         keep_dims = False
-        value = torch.reshape(
-            value,
-            (
-                value.shape[0],
-                math.ceil(value.shape[1] / args.group_size),
-                args.group_size,
-            ),
+
+        reshaped_dims = (
+            math.ceil(value.shape[-1] / args.group_size),
+            args.group_size,
         )
+        value = value.unflatten(-1, reshaped_dims)
+
     else:
         supported_strategies = (
             QuantizationStrategy.TOKEN,

src/compressed_tensors/transform/apply.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import torch
+from compressed_tensors import TRANSFORM_CONFIG_NAME
 from compressed_tensors.transform import TransformConfig, TransformFactory
 
 
@@ -30,3 +31,6 @@ def apply_transform_config(model: torch.nn.Module, config: TransformConfig):
     for name, scheme in config.config_groups.items():
         factory = TransformFactory.from_scheme(scheme, name=name)
         factory.apply_to_model(model)
+
+    # attach config to model for compression/serialization
+    setattr(model, TRANSFORM_CONFIG_NAME, config)

src/compressed_tensors/transform/factory/base.py

@@ -14,11 +14,11 @@
 
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Set
 
 import torch
 import torch.nn.utils.parametrize as P
-from compressed_tensors import InternalModule, match_named_modules
+from compressed_tensors.utils.internal import InternalModule
 from compressed_tensors.registry.registry import RegistryMixin, T
 from compressed_tensors.transform import (
     TransformArgs,
@@ -164,10 +164,6 @@ def _update_tied_weights(self):
         which is used by transformers to detect and remove shared pointers
         during saving
         """
-        # avoid issues with this method being called twice
-        for transform in self.transforms:
-            transform._dynamic_tied_weights_keys = list()
-
         # map from data_ptrs to keys
         ptr_to_keys: dict[int, List[Tuple[TransformBase, str]]] = defaultdict(list)
         for transform in self.transforms:
@@ -184,7 +180,7 @@ def _update_tied_weights(self):
                 tensor = getattr(shared_keys[0][0], shared_keys[0][1])
 
                 for transform, name in shared_keys:
-                    transform._dynamic_tied_weights_keys.append(name)
+                    transform._dynamic_tied_weights_keys.add(name)
                     setattr(transform, name, tensor)
 
 
@@ -195,11 +191,11 @@ class TransformBase(InternalModule, ABC):
 
     args: TransformArgs
     weight: Parameter
-    _dynamic_tied_weights_keys: List[str]
+    _dynamic_tied_weights_keys: Set[str]
 
     def __init__(self):
         super().__init__()
-        self._dynamic_tied_weights_keys = list()
+        self._dynamic_tied_weights_keys = set()
 
     @abstractmethod
     def forward(self, value: Tensor) -> Tensor:

src/compressed_tensors/transform/factory/hadamard.py

@@ -13,7 +13,11 @@
 # limitations under the License.
 
 import math
+<<<<<<< HEAD
 from typing import Optional
+=======
+from typing import Optional, Union
+>>>>>>> kylesayrs/serialize-tconfig
 
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme