populate _dynamic_tied_weights_keys

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

Author: kylesayrs

src/compressed_tensors/transform/factory/base.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from collections import defaultdict
+from typing import List, Optional, Tuple
 
 import torch
 import torch.nn.utils.parametrize as P
@@ -48,10 +49,13 @@ class TransformFactory(RegistryMixin, ABC):
     :param seed: random seed used to transform weight randomization
     """
 
+    transforms: List["TransformBase"]
+
     def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
         self.name = name
         self.scheme = scheme
         self.generator = torch.Generator()
+        self.transforms = list()
         if seed is not None:
             self.generator.manual_seed(seed)
 
@@ -90,6 +94,8 @@ def apply_to_model(self, model: Module):
                 if is_target(name, module, arg.targets, arg.ignore):
                     self._apply_to_module(module, arg)
 
+        self._update_tied_weights()
+
     def _apply_to_module(self, module: Module, args: TransformArgs):
         """
         Create transforms and apply them to the module
@@ -145,6 +151,28 @@ def output_hook(_, _input, output):
         else:
             raise NotImplementedError()
 
+    def _update_tied_weights(self):
+        """
+        Populate the `_dynamic_tied_weights_keys` attribute of transforms,
+        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:
+            for name, param in transform.named_parameters(recurse=False):
+                ptr_to_keys[param.data_ptr()].append((transform, name))
+
+        # populate `_dynamic_tied_weights_keys` if there is more than one key
+        for shared_keys in ptr_to_keys.values():
+            if len(shared_keys) > 1:
+                for transform, name in shared_keys:
+                    transform._dynamic_tied_weights_keys.append(name)
+
 
 class TransformBase(Module, ABC):
     """
@@ -153,6 +181,11 @@ class TransformBase(Module, ABC):
 
     args: TransformArgs
     weight: Parameter
+    _dynamic_tied_weights_keys: List[str]
+
+    def __init__(self):
+        super().__init__()
+        self._dynamic_tied_weights_keys = list()
 
     @abstractmethod
     def forward(self, value: Tensor) -> Tensor:

src/compressed_tensors/transform/factory/hadamard.py

@@ -60,7 +60,10 @@ def create_transform(self, module: Module, args: TransformArgs):
         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)
+
+        transform = HadamardTransform(weight, perm, args)
+        self.transforms.append(transform)
+        return transform
 
     def _create_weight(
         self,

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -59,7 +59,9 @@ def create_transform(self, module: Module, args: TransformArgs):
         if args.inverse:
             weight = self.inverses[weight]
 
-        return RandomMatrixTransform(weight, args)
+        transform = RandomMatrixTransform(weight, args)
+        self.transforms.append(transform)
+        return transform
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
         # TODO: verify that weight is invertible (has non-zero determinant)