move loading step to transformers

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -39,7 +39,6 @@
     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
@@ -472,8 +471,6 @@ def decompress(self, model_path: str, model: Module):
                 )
                 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
             )

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -84,11 +84,16 @@ def load_transforms(model: Module, model_name_or_path: str):
 
     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():
+            for transform_name, transform_values 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)
+                transform = transform_values.get("transform")
+                transform.register_to_module(name=transform_name, module=submodule)
+                transform.update_transform(
+                    module=submodule, data=transform_data, name=transform_name
+                )
 
 
 def load_pretrained_quantization(model: Module, model_name_or_path: str):
@@ -194,7 +199,9 @@ def process_transforms_config(
                             dtype=dtype,
                             **transform_creation_args,
                         )
-                    transform.register_to_module(name=transform_name, module=submodule)
+                        transform.register_to_module(
+                            name=transform_name, module=submodule
+                        )
 
                     # add relevant transform data to the submodule as well
                     data = {