key by weight only

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

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -42,6 +42,7 @@ def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = Non
         super().__init__(name, scheme, seed)
         self.weights = ParameterizedDefaultDict(self._create_weight)
         self.perms = ParameterizedDefaultDict(self._create_permutation)
+        self._shared_tensors_device = None
 
     def create_transform(self, module: Module, args: TransformArgs):
         """
@@ -57,20 +58,33 @@ def create_transform(self, module: Module, args: TransformArgs):
         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)
+        factory_kwargs = {"device": device, "construct_device": exec_device}
+        weight = self.weights.get(size, factory_kwargs=factory_kwargs)
+        # TODO: permutations should be keyed by fused modules, not weight
         perm = self.perms[weight] if self.scheme.randomize else None
         return HadamardTransform(weight, perm, self.scheme, args, type(module))
 
     def _create_weight(
         self,
         size: int,
-        dtype: dtype,
         device: device,
         construct_device: device,
     ) -> Parameter:
-        # construct on execution device, cache on offload device
-        data = deterministic_hadamard_matrix(size, dtype, construct_device)
+        # check that shared tensors device is consistent
+        if self._shared_tensors_device is None:
+            self._shared_tensors_device = device
+
+        if device != self._shared_tensors_device:
+            raise NotImplementedError(
+                "Creating multi-gpu transform weights are not supported as of now due "
+                "to the limitations of shared tensors across GPUs."
+                # in the future, tensors can be shared within GPUs,
+                # and can be all-reduced during updates and compression
+            )
+
+        # construct on execution device, cache shared tensor on offload device
+        precision = self.scheme.precision
+        data = deterministic_hadamard_matrix(size, precision, construct_device)
         data = data.to(device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -41,6 +41,7 @@ def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = Non
         super().__init__(name, scheme, seed)
         self.weights = ParameterizedDefaultDict(self._create_weight)
         self.inverses = ParameterizedDefaultDict(self._create_inverse)
+        self._shared_tensors_device = None
 
     def create_transform(self, module: Module, args: TransformArgs):
         """
@@ -52,19 +53,34 @@ def create_transform(self, module: Module, args: TransformArgs):
         """
         assert hasattr(module, "weight")
         size = get_transform_size(module, args.location, self.scheme.head_dim)
-        dtype = self.scheme.precision
         device = get_offloaded_device(module)
 
-        weight = self.weights[size, dtype, device]
+        factory_kwargs = {"device": device}
+        weight = self.weights.get(size, factory_kwargs=factory_kwargs)
         if args.inverse:
             weight = self.inverses[weight]
 
         return RandomMatrixTransform(weight, self.scheme, args, type(module))
 
-    def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
-        # TODO: verify that weight is invertible (has non-zero determinant)
+    def _create_weight(self, size: int, device: device) -> Parameter:
+        # check that shared tensors device is consistent
+        if self._shared_tensors_device is None:
+            self._shared_tensors_device = device
+
+        if device != self._shared_tensors_device:
+            raise NotImplementedError(
+                "Creating multi-gpu transform weights are not supported as of now due "
+                "to the limitations of shared tensors across GPUs"
+                # in the future, tensors can be shared within GPUs,
+                # and can be all-reduced during updates and compression
+            )
+
+        # TODO: construct such that weight is invertible (has non-zero determinant)
         data = torch.rand(
-            (size, size), generator=self.generator, dtype=dtype, device=device
+            (size, size),
+            generator=self.generator,
+            dtype=self.scheme.precision,
+            device=device,
         )
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 

src/compressed_tensors/transform/factory/random_hadamard.py

@@ -31,11 +31,11 @@ class RandomHadamardFactory(HadamardFactory):
     def _create_weight(
         self,
         size: int,
-        dtype: dtype,
         device: device,
         construct_device: device,
     ) -> Parameter:
         # construct on execution device, cache on offload device
-        data = random_hadamard_matrix(size, dtype, construct_device, self.generator)
+        precision = self.scheme.precision
+        data = random_hadamard_matrix(size, precision, construct_device, self.generator)
         data = data.to(device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)