save construction device changes for later

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

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -22,7 +22,7 @@
     apply_transform_weight,
     get_matrix_size,
 )
-from compressed_tensors.utils import get_execution_device, get_offloaded_device
+from compressed_tensors.utils import get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
 from torch import Tensor, device, dtype
 from torch.nn import Linear, Module, Parameter
@@ -41,7 +41,6 @@ class HadamardFactory(TransformFactory):
     def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
         super().__init__(name, scheme, seed)
         self.weights = ParameterizedDefaultDict(self._create_weight)
-        self._exec_device = torch.device("cpu")
 
     def create_transform(self, module: Module, args: TransformArgs):
         """
@@ -55,21 +54,13 @@ def create_transform(self, module: Module, args: TransformArgs):
         size = get_matrix_size(module, args.location)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
-        exec_device = get_execution_device(module)
 
-        weight = self.weights.get(size, dtype, device, construct_device=exec_device)
+        weight = self.weights[size, dtype, device]
         return HadamardTransform(weight, args)
 
-    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)
-        data = data.to(device=device)
+    def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        data = deterministic_hadamard_matrix(size, dtype, device)
+        data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 
 

src/compressed_tensors/transform/factory/random_hadamard.py

@@ -28,14 +28,7 @@ class RandomHadamardFactory(HadamardFactory):
     :param seed: random seed used to transform weight randomization
     """
 
-    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)
-        data = data.to(device=device)
+    def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        data = random_hadamard_matrix(size, dtype, device, self.generator)
+        data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)

src/compressed_tensors/utils/helpers.py

@@ -373,16 +373,11 @@ class ParameterizedDefaultDict(dict):
 
     def __init__(self, default_factory: Callable[[Any], Any]):
         self.default_factory = default_factory
-        self._kwargs = {}
 
-    def __missing__(self, key: Any) -> Any:
+    def __missing__(self, key):
         if isinstance(key, tuple):
-            value = self.default_factory(*key, **self._kwargs)
+            value = self.default_factory(*key)
         else:
-            value = self.default_factory(key, **self._kwargs)
+            value = self.default_factory(key)
         self[key] = value
         return value
-
-    def get(self, *args, **kwargs) -> Any:
-        with patch_attr(self, "_kwargs", kwargs):
-            return self[args]