eagerly generate with precision

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -24,7 +24,7 @@
 )
 from compressed_tensors.utils import get_execution_device, get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
-from torch import Tensor, device
+from torch import Tensor, device, dtype
 from torch.nn import Module, Parameter
 
 
@@ -55,8 +55,13 @@ def create_transform(self, module: Module, args: TransformArgs):
         size = get_transform_size(module, args.location, self.scheme.head_dim)
         exec_device = get_execution_device(module)
         device = get_offloaded_device(module)
+        precision = self.scheme.precision if args.is_online() else torch.float64
 
-        factory_kwargs = {"device": device, "construct_device": exec_device}
+        factory_kwargs = {
+            "device": device,
+            "construct_device": exec_device,
+            "precision": precision,
+        }
         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
@@ -67,8 +72,8 @@ def _create_weight(
         size: int,
         device: device,
         construct_device: device,
+        precision: dtype,
     ) -> Parameter:
-        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)
@@ -93,8 +98,7 @@ def __init__(
         self.scheme = scheme
         self.args = args
         self.module_type = module_type
-        self._scale = torch.tensor(weight.size(0), dtype=self.scheme.precision).sqrt()
-        self._precision = scheme.precision if args.is_online() else torch.float64
+        self._scale = torch.tensor(weight.size(0), dtype=torch.float64).sqrt()
 
     def forward(self, value: Tensor) -> Tensor:
         weight = self.weight
@@ -107,8 +111,8 @@ def forward(self, value: Tensor) -> Tensor:
 
         return (
             apply_transform_weight(
-                weight.to(dtype=self._precision, device=value.device),
-                value.to(self._precision),
+                weight.to(device=value.device),
+                value.to(dtype=weight.dtype),
                 self.args.location,
                 self.module_type,
             )

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -23,7 +23,7 @@
 )
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
 from compressed_tensors.utils.offload import get_offloaded_device
-from torch import Tensor, device
+from torch import Tensor, device, dtype
 from torch.nn import Module, Parameter
 
 
@@ -53,20 +53,21 @@ def create_transform(self, module: Module, args: TransformArgs):
         assert hasattr(module, "weight")
         size = get_transform_size(module, args.location, self.scheme.head_dim)
         device = get_offloaded_device(module)
+        precision = self.scheme.precision if args.is_online() else torch.float64
 
-        factory_kwargs = {"device": device}
+        factory_kwargs = {"device": device, "precision": precision}
         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, device: device) -> Parameter:
+    def _create_weight(self, size: int, device: device, precision: dtype) -> Parameter:
         # TODO: construct such that weight is invertible (has non-zero determinant)
         data = torch.rand(
             (size, size),
             generator=self.generator,
-            dtype=self.scheme.precision,
+            dtype=precision,
             device=device,
         )
         return Parameter(data, requires_grad=self.scheme.requires_grad)
@@ -90,21 +91,20 @@ def __init__(
         self.scheme = scheme
         self.args = args
         self.module_type = module_type
-        self._precision = scheme.precision if args.is_online() else torch.float64
 
     def forward(self, value: Tensor) -> Parameter:
         return apply_transform_weight(
-            self.weight.to(dtype=self._precision, device=value.device),
-            value.to(self._precision),
+            self.weight.to(device=value.device),
+            value.to(dtype=self.weight.dtype),
             self.args.location,
             self.module_type,
         ).to(value.dtype)
 
     def right_inverse(self, value: Tensor) -> Tensor:
         inverse = high_precision_invert(self.weight)
         return apply_transform_weight(
-            inverse.to(dtype=self._precision, device=value.device),
-            value.to(self._precision),
+            inverse.to(device=value.device),
+            value.to(dtype=inverse.dtype),
             self.args.location,
             self.module_type,
         ).to(value.dtype)

src/compressed_tensors/transform/factory/random_hadamard.py

@@ -14,7 +14,7 @@
 
 from compressed_tensors.transform import HadamardFactory, TransformFactory
 from compressed_tensors.transform.utils.hadamard import random_hadamard_matrix
-from torch import device
+from torch import device, dtype
 from torch.nn import Parameter
 
 
@@ -33,8 +33,8 @@ def _create_weight(
         size: int,
         device: device,
         construct_device: device,
+        precision: dtype,
     ) -> Parameter:
-        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)