undo transform precision

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

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import Optional
 
-import math
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
 from compressed_tensors.transform.factory.base import TransformBase, TransformFactory
@@ -104,6 +104,7 @@ def forward(self, value: Tensor) -> Tensor:
         if self.args.inverse:
             weight = weight.T
 
-        return apply_transform_weight(
-            weight, value, self.args.location, self.module_type
-        ) / self._scale
+        return (
+            apply_transform_weight(weight, value, self.args.location, self.module_type)
+            / self._scale
+        )

src/compressed_tensors/transform/utils/matrix.py

@@ -21,9 +21,6 @@
 __all__ = ["get_transform_size", "apply_transform_weight"]
 
 
-TRANSFORM_PRECISION = torch.float64
-
-
 def get_transform_size(
     module: torch.nn.Module,
     location: TransformLocation,
@@ -88,14 +85,8 @@ def apply_transform_weight(
     num_heads = value.shape[axis] // head_dim
     value = value.unflatten(axis, (num_heads, head_dim))
 
-    # cast to transform precision
-    value_dtype = value.dtype
-
     # apply transform
-    value = fn(weight.to(TRANSFORM_PRECISION), value.to(TRANSFORM_PRECISION))
-
-    # [undo] cast to transform precision
-    value = value.to(value_dtype)
+    value = fn(weight, value)
 
     # [undo] reshape for head_dim
     value = value.flatten(axis - 1, axis)

src/compressed_tensors/utils/__init__.py

@@ -15,9 +15,9 @@
 
 from .helpers import *
 from .internal import *
+from .match import *
 from .offload import *
 from .permutations_24 import *
 from .permute import *
 from .safetensors_load import *
 from .semi_structured_conversions import *
-from .match import *

src/compressed_tensors/utils/match.py

@@ -1,9 +1,24 @@
-from typing import Iterable, Tuple
-from collections.abc import Generator
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
+import logging
 import re
+from collections.abc import Generator
+from typing import Iterable, Tuple
+
 import torch
-import logging
+
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -15,14 +30,14 @@ def match_named_modules(
     model: torch.nn.Module,
     targets: Iterable[str] = tuple(),
     ignore: Iterable[str] = tuple(),
-    warn_on_fail: bool = True
+    warn_on_fail: bool = True,
 ) -> Generator[Tuple[str, torch.nn.Module], None, None]:
     unmatched_targets = set(targets)
     for name, module in model.named_modules():
         for target in targets:
             if is_match(name, module, target):
                 unmatched_targets.remove(target)
-        
+
                 if not any(is_match(name, module, ign) for ign in ignore):
                     yield name, module
 
@@ -32,6 +47,7 @@ def match_named_modules(
                 f"Could not match `{target}` in instance of {model.__class__.__name__}"
             )
 
+
 def is_match(name: str, module: torch.nn.Module, target: str) -> bool:
     return _match_name(name, target) or _match_class(module, target)
 
@@ -46,7 +62,7 @@ def _match_name(name: str, target: str) -> bool:
 def _match_class(module: torch.nn.Module, target: str) -> bool:
     """
     Will never match against a regex pattern since `:` is not allowed in class names
-    
+
     """
     return any(
         issubclass(cls, torch.nn.Module) and cls.__name__ == target