ensure serializable

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/transform/factory/base.py

@@ -103,9 +103,17 @@ def _apply_to_module(self, module: Module, args: TransformArgs):
         :param module: target module to apply transforms to
         :param args: defines how the transform will be applied to the target module
         """
+        if has_offloaded_params(module):
+            if module._hf_hook.place_submodules:
+                raise NotImplementedError(
+                    "Applying transforms to offloaded submodules with "
+                    "`place_submodules=True` is not supported"
+                )
+
         # create transform as submodule
         transform_name = f"{self.name}_{args.location.value}"
         transform = self.create_transform(module, args)
+        self.transforms.append(transform)
         register_offload_module(module, transform_name, transform)
 
         # register input transformation hook
@@ -136,8 +144,9 @@ def input_hook(_, args):
                     raise ValueError("Offloaded training is not supported")
                 P.register_parametrization(module, "weight", transform)
 
-            # transform is no longer needed (unfusing is not supported)
-            delete_offload_module(module, transform_name)
+            else:
+                # transform is no longer needed (unfusing is not supported)
+                delete_offload_module(module, transform_name)
 
         # register output transformation hook
         elif args.location == TransformLocation.OUTPUT:
@@ -165,13 +174,20 @@ def _update_tied_weights(self):
         ptr_to_keys: dict[int, List[Tuple[TransformBase, str]]] = defaultdict(list)
         for transform in self.transforms:
             for name, param in transform.named_parameters(recurse=False):
+                # NOTE: previously asserted that parent._hf_hook.place_submodules=False
+                if has_offloaded_params(transform):
+                    param = transform._hf_hook.weights_map[name]
                 ptr_to_keys[param.data_ptr()].append((transform, name))
 
         # populate `_dynamic_tied_weights_keys` if there is more than one key
+        # and ensure that they share tensors
         for shared_keys in ptr_to_keys.values():
             if len(shared_keys) > 1:
+                tensor = getattr(shared_keys[0][0], shared_keys[0][1])
+
                 for transform, name in shared_keys:
                     transform._dynamic_tied_weights_keys.append(name)
+                    setattr(transform, name, tensor)
 
 
 class TransformBase(Module, ABC):

src/compressed_tensors/transform/factory/hadamard.py

@@ -61,9 +61,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         weight = self.weights.get(size, dtype, device, factory_kwargs=factory_kwargs)
         perm = self.perms[weight] if self.scheme.randomize else None
 
-        transform = HadamardTransform(weight, perm, args)
-        self.transforms.append(transform)
-        return transform
+        return HadamardTransform(weight, perm, args)
 
     def _create_weight(
         self,

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -59,9 +59,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         if args.inverse:
             weight = self.inverses[weight]
 
-        transform = RandomMatrixTransform(weight, args)
-        self.transforms.append(transform)
-        return transform
+        return RandomMatrixTransform(weight, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
         # TODO: verify that weight is invertible (has non-zero determinant)
@@ -72,6 +70,7 @@ def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
 
     def _create_inverse(self, weight: Parameter) -> Parameter:
         data = high_precision_invert(weight.data)
+        data = data.contiguous()  # ensure proper serialization
         return Parameter(data, requires_grad=False)
 
 

tests/test_transform/conftest.py

@@ -14,12 +14,13 @@
 
 import pytest
 import torch
-from compressed_tensors.transform import TransformArgs
+from compressed_tensors.transform import TransformArgs, TransformFactory
+from transformers import PretrainedConfig, PreTrainedModel
 
 
-class TransformableModel(torch.nn.Module):
+class TransformableModel(PreTrainedModel):
     def __init__(self, *sizes):
-        super().__init__()
+        super().__init__(config=PretrainedConfig())
         self.fcs = torch.nn.ModuleList(
             [
                 torch.nn.Linear(sizes[index], sizes[index + 1], bias=False)

tests/test_transform/factory/test_correctness.py

@@ -22,6 +22,7 @@
     apply_transform_config,
 )
 from compressed_tensors.utils import offloaded_dispatch
+from tests.test_transform.conftest import scheme_kwargs
 from tests.testing_utils import requires_accelerate, requires_gpu
 
 

tests/test_transform/factory/test_memory.py

@@ -20,8 +20,8 @@
     TransformArgs,
     TransformBase,
     TransformConfig,
-    TransformFactory,
     TransformScheme,
+    TransformFactory,
     apply_transform_config,
 )
 from compressed_tensors.utils import align_modules, offloaded_dispatch

tests/test_transform/factory/test_serialization.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2021 - present / Neuralmagic, 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 pytest
+import torch
+from compressed_tensors.transform import (
+    TransformConfig,
+    TransformScheme,
+    apply_transform_config,
+)
+from compressed_tensors.utils import offloaded_dispatch
+from tests.test_transform.conftest import scheme_kwargs
+from tests.testing_utils import requires_accelerate, requires_gpu
+
+
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_serialization(scheme_kwargs, model_apply, tmp_path, offload=False):
+    # get model, maybe offload
+    model, apply = model_apply
+    if offload:
+        offloaded_dispatch(model, torch.device("cuda"))
+
+    # apply transforms to model
+    config = TransformConfig(
+        config_groups={"": TransformScheme(**scheme_kwargs, apply=apply)}
+    )
+    apply_transform_config(model, config)
+
+    # save model
+    model.save_pretrained(tmp_path)
+
+    # TODO: reload model
+
+
+@requires_gpu
+@requires_accelerate()
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_serialization_offload(scheme_kwargs, model_apply, tmp_path):
+    test_serialization(scheme_kwargs, model_apply, tmp_path, offload=True)