cleanup, construct on dtype, change default

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

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -53,15 +53,14 @@ def create_transform(self, module: Module, args: TransformArgs):
         """
         assert hasattr(module, "weight")
         size = get_transform_size(module, args.location, self.scheme.head_dim)
-        dtype = module.weight.dtype
+        dtype = self.scheme.precision
         device = get_offloaded_device(module)
         exec_device = get_execution_device(module)
-        precision = self.scheme.precision
 
         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, precision, type(module))
+        return HadamardTransform(weight, perm, self.scheme, args, type(module))
 
     def _create_weight(
         self,
@@ -85,17 +84,17 @@ def __init__(
         self,
         weight: Parameter,
         perm: Optional[Parameter],
+        scheme: TransformScheme,
         args: TransformArgs,
-        precision: torch.dtype,
         module_type: type[torch.nn.Module],
     ):
         super().__init__()
         self.weight = weight
         self.perm = perm
+        self.scheme = scheme
         self.args = args
-        self.precision = precision
         self.module_type = module_type
-        self._scale = torch.tensor(weight.size(0), dtype=self.precision).sqrt()
+        self._scale = torch.tensor(weight.size(0), dtype=self.scheme.precision).sqrt()
 
     def forward(self, value: Tensor) -> Tensor:
         weight = self.weight
@@ -108,8 +107,8 @@ def forward(self, value: Tensor) -> Tensor:
 
         return (
             apply_transform_weight(
-                weight.to(self.precision),
-                value.to(self.precision),
+                weight.to(self.scheme.precision),
+                value.to(self.scheme.precision),
                 self.args.location,
                 self.module_type,
             )

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -24,7 +24,7 @@
 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
+from torch.nn import Module, Parameter
 
 
 @TransformFactory.register("random-matrix")
@@ -52,7 +52,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         """
         assert hasattr(module, "weight")
         size = get_transform_size(module, args.location, self.scheme.head_dim)
-        dtype = module.weight.dtype
+        dtype = self.scheme.precision
         device = get_offloaded_device(module)
         precision = self.scheme.precision
 
@@ -78,29 +78,29 @@ class RandomMatrixTransform(TransformBase):
     def __init__(
         self,
         weight: Tensor,
+        scheme: TransformScheme,
         args: TransformArgs,
-        precision: torch.dtype,
         module_type: type[torch.nn.Module],
     ):
         super().__init__()
         self.weight = weight  # is an inverse if args.inverse
+        self.scheme = scheme
         self.args = args
-        self.precision = precision
         self.module_type = module_type
 
     def forward(self, value: Tensor) -> Parameter:
         return apply_transform_weight(
-            self.weight.to(self.precision),
-            value.to(self.precision),
+            self.weight.to(self.scheme.precision),
+            value.to(self.scheme.precision),
             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(self.precision),
-            value.to(self.precision),
+            inverse.to(self.scheme.precision),
+            value.to(self.scheme.precision),
             self.args.location,
             self.module_type,
         ).to(value.dtype)

src/compressed_tensors/transform/transform_scheme.py

@@ -36,11 +36,13 @@ class TransformScheme(BaseModel):
     :param randomize: True if uniquely randomized transform weights should be used,
         otherwise use identical transform weights where applicable
     :param requires_grad: True if weights include gradients for training
+    :param precision: Precision at which this transform should be applied. This applies
+        to both weight fusing and online rotations
     """
 
     type: str
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
     head_dim: Optional[int] = Field(default=None)
-    precision: TorchDtype = Field(default=torch.bfloat16)
+    precision: TorchDtype = Field(default=torch.float32)