Merge branch 'main' into support_multi_compressor

Author: dsikka

pyproject.toml

@@ -5,3 +5,12 @@ build-backend = "setuptools.build_meta"
 [tool.black]
 line-length = 88
 target-version = ['py36']
+
+[tool.pytest.ini_options]
+markers = [
+    "unit: tests to ensure code correctness and regression test functionality",
+    "smoke: quick tests to check basic functionality",
+    "sanity: tests to ensure that new changes do not break existing functionality",
+    "regression: detailed tests to ensure major functions work correctly",
+    "integration: tests which integrate with a third party service such as HF",
+]

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)
+        # TODO: transform config is not support by CompressedTensorsConfig yet
+
         if sparsity_config is None and quantization_config is None:
             return None
 
@@ -177,7 +182,6 @@ def from_pretrained_model(
             algorithm
         :return: compressor for the configs, or None if model is not compressed
         """
-
         if quantization_format is not None:
             # llmcompressor incorrectly passes in a CompressionFormat when
             # the value string is expected - handle both cases
@@ -194,17 +198,23 @@ def from_pretrained_model(
         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,
+            transform_config=transform_config,
         )
 
     @staticmethod
@@ -283,13 +293,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[
             Dict[str, Union[BaseQuantizationCompressor, DenseCompressor]]
         ] = None
+        # no transform compressor is required
 
         if sparsity_config is not None:
             self.sparsity_compressor = BaseCompressor.load_from_registry(
@@ -718,44 +732,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:
-            return
-
-        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."
-            )
+        # 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
 
-        with open(config_file_path, "r") as config_file:
-            config_data = json.load(config_file)
-
-        # required metadata whenever a quantization or sparsity config is present
+        # write to config.json file, regardless of whether it exists already
         # 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
+        config_file_path = os.path.join(save_directory, CONFIG_NAME)
+        if os.path.exists(config_file_path):
+            with open(config_file_path, "r") as file:
+                config_data = json.load(file)
         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
+            config_data = {}
+
+        # 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 {}
+        )
 
+        # 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/quantization/quant_args.py

@@ -19,7 +19,7 @@
 import torch
 from compressed_tensors.utils import Aliasable
 from compressed_tensors.utils.helpers import deprecated
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 
 __all__ = [
@@ -358,6 +358,8 @@ def pytorch_dtype(self) -> torch.dtype:
     def get_observer(self) -> str:
         return self.observer
 
+    model_config = ConfigDict(extra="forbid")
+
 
 def round_to_quantized_type(
     tensor: torch.Tensor, args: QuantizationArgs

src/compressed_tensors/quantization/quant_config.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Dict, List, Optional, Union
+from typing import Annotated, Any, Dict, List, Optional, Union
 
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.quant_args import DynamicType, QuantizationArgs
@@ -26,7 +26,7 @@
     module_type,
     parse_out_kv_cache_args,
 )
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 from torch.nn import Module
 
 
@@ -142,6 +142,9 @@ class QuantizationConfig(BaseModel):
     quantization_status: QuantizationStatus = QuantizationStatus.INITIALIZED
     global_compression_ratio: Optional[float] = None
     ignore: Optional[List[str]] = Field(default_factory=list)
+    # `run_compressed` is a dummy, unused arg for backwards compatibility
+    # see: https://github.com/huggingface/transformers/pull/39324
+    run_compressed: Annotated[Any, Field(exclude=True)] = None
 
     def model_post_init(self, __context):
         """
@@ -254,3 +257,6 @@ def requires_calibration_data(self):
                     return True
 
         return False
+
+    # TODO set `extra="forbid"` when upstream transformers is compatible
+    model_config = ConfigDict(extra="ignore")

src/compressed_tensors/quantization/quant_scheme.py

@@ -14,7 +14,7 @@
 
 import warnings
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import List, Optional
 
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.quant_args import (
@@ -23,7 +23,7 @@
     QuantizationStrategy,
     QuantizationType,
 )
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, ConfigDict, model_validator
 
 
 __all__ = [
@@ -83,6 +83,8 @@ def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":
 
         return model
 
+    model_config = ConfigDict(extra="forbid")
+
 
 """
 Pre-Set Quantization Scheme Args

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

@@ -18,7 +18,6 @@
 
 import torch
 import torch.nn.utils.parametrize as P
-from compressed_tensors import InternalModule
 from compressed_tensors.registry.registry import RegistryMixin, T
 from compressed_tensors.transform import (
     TransformArgs,
@@ -34,6 +33,7 @@
     register_offload_module,
     update_offload_parameter,
 )
+from compressed_tensors.utils.internal import InternalModule
 from torch import Tensor
 from torch.nn import Module, Parameter
 

src/compressed_tensors/transform/factory/hadamard.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
@@ -26,7 +25,7 @@
 from compressed_tensors.utils import get_execution_device, get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
 from torch import Tensor, device, dtype
-from torch.nn import Linear, Module, Parameter
+from torch.nn import Module, Parameter
 
 
 @TransformFactory.register("hadamard")
@@ -54,14 +53,14 @@ def create_transform(self, module: Module, args: TransformArgs):
         """
         assert hasattr(module, "weight")
         size = get_transform_size(module, args.location, self.scheme.head_dim)
-        dtype = module.weight.dtype
+        dtype = self.scheme.precision
         device = get_offloaded_device(module)
         exec_device = get_execution_device(module)
 
         factory_kwargs = {"construct_device": exec_device}
         weight = self.weights.get(size, dtype, device, factory_kwargs=factory_kwargs)
         perm = self.perms[weight] if self.scheme.randomize else None
-        return HadamardTransform(weight, perm, args, type(module))
+        return HadamardTransform(weight, perm, self.scheme, args, type(module))
 
     def _create_weight(
         self,
@@ -85,15 +84,18 @@ def __init__(
         self,
         weight: Parameter,
         perm: Optional[Parameter],
+        scheme: TransformScheme,
         args: TransformArgs,
         module_type: type[torch.nn.Module],
     ):
         super().__init__()
         self.weight = weight
         self.perm = perm
+        self.scheme = scheme
         self.args = args
         self.module_type = module_type
-        self._scale = math.sqrt(weight.size(0))
+        self._scale = torch.tensor(weight.size(0), dtype=self.scheme.precision).sqrt()
+        self._precision = scheme.precision if args.is_online() else torch.float64
 
     def forward(self, value: Tensor) -> Tensor:
         weight = self.weight
@@ -105,6 +107,11 @@ def forward(self, value: Tensor) -> Tensor:
             weight = weight.T
 
         return (
-            apply_transform_weight(weight, value, self.args.location, self.module_type)
+            apply_transform_weight(
+                weight.to(self._precision),
+                value.to(self._precision),
+                self.args.location,
+                self.module_type,
+            )
             / self._scale
-        )
+        ).to(value.dtype)