Merge branch 'kylesayrs/transform_construct_cache_device' into kylesayrs/transform_apply

Author: kylesayrs

setup.py

@@ -88,7 +88,7 @@ def _setup_packages() -> List:
     )
 
 def _setup_install_requires() -> List:
-    return ["torch>=1.7.0", "transformers", "pydantic>=2.0"]
+    return ["torch>=1.7.0", "transformers", "pydantic>=2.0", "frozendict"]
 
 def _setup_extras() -> Dict:
     return {

src/compressed_tensors/transform/factory/hadamard.py

@@ -57,7 +57,8 @@ def create_transform(self, module: Module, args: TransformArgs):
         device = get_offloaded_device(module)
         exec_device = get_execution_device(module)
 
-        weight = self.weights.get(size, dtype, device, construct_device=exec_device)
+        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)
 

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -62,6 +62,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         return RandomMatrixTransform(weight, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        # TODO: verify that weight is invertable (has non-zero determinant)
         data = torch.rand(
             (size, size), generator=self.generator, dtype=dtype, device=device
         )

src/compressed_tensors/transform/factory/random_hadamard.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
 from compressed_tensors.transform import HadamardFactory, TransformFactory
 from compressed_tensors.transform.utils.hadamard import random_hadamard_matrix
 from torch import device, dtype

src/compressed_tensors/utils/helpers.py

@@ -15,10 +15,11 @@
 import contextlib
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional
 
 import numpy
 import torch
+from frozendict import frozendict
 from transformers import AutoConfig
 
 
@@ -373,16 +374,23 @@ class ParameterizedDefaultDict(dict):
 
     def __init__(self, default_factory: Callable[[Any], Any]):
         self.default_factory = default_factory
-        self._kwargs = {}
+        self._factory_kwargs = frozendict()
 
     def __missing__(self, key: Any) -> Any:
         if isinstance(key, tuple):
-            value = self.default_factory(*key, **self._kwargs)
+            value = self.default_factory(*key, **self._factory_kwargs)
         else:
-            value = self.default_factory(key, **self._kwargs)
+            value = self.default_factory(key, **self._factory_kwargs)
         self[key] = value
         return value
 
-    def get(self, *args, **kwargs) -> Any:
-        with patch_attr(self, "_kwargs", kwargs):
+    def get(self, *args, factory_kwargs: Mapping = frozendict()) -> Any:
+        """
+        Similar to `__getitem__`, but allows passing kwargs to factory function
+
+        :param \\*args: args whose tuple will value will be treated as key
+        :param factory_kwargs: keyword arguments to pass to `default_factory`
+        :return: dictionary entry for given key
+        """
+        with patch_attr(self, "_factory_kwargs", factory_kwargs):
             return self[args]