[Transform] Construct on GPU, cache on CPU (#352)

kylesayrs · web-flow · commit 6c63987a2d2d · 2025-07-07T12:48:24.000-04:00
* use hadamards database file Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * try manifest Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * try setup, update hadamards list Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * fix setup Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * add docstrings, cleanup Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * fix setup, thank you @dbarbuzzi Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * remove numpy, add tests Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * solidify dtype, add gpu tests Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * fix docstring Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * add device option Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * construct on execution device, cache on offload device Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * save construction device changes for later Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * construct on execution device, cache on offload device * cite nja sloane Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * remove dreg Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * put on device via safe_open Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * nits and docstrings Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * update docstring Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * construct with same dtype, constructing on fp32 found no difference Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * remove unnecessary imports Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * use factory_kwargs Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * add frozen dict to deps Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * correct typo Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> * fix missing import Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> --------- Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
diff --git a/setup.py b/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 {
diff --git a/src/compressed_tensors/transform/factory/hadamard.py b/src/compressed_tensors/transform/factory/hadamard.py
@@ -22,7 +22,7 @@
     apply_transform_weight,
     get_matrix_size,
 )
-from compressed_tensors.utils import get_offloaded_device
+from compressed_tensors.utils import get_execution_device, get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
 from torch import Tensor, device, dtype
 from torch.nn import Linear, Module, Parameter
@@ -55,14 +55,23 @@ 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[size, dtype, 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)
 
-    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)
+    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)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 
     def _create_permutation(self, weight: Parameter) -> Parameter:
diff --git a/src/compressed_tensors/transform/factory/matrix_multiply.py b/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 invertible (has non-zero determinant)
         data = torch.rand(
             (size, size), generator=self.generator, dtype=dtype, device=device
         )
diff --git a/src/compressed_tensors/transform/factory/random_hadamard.py b/src/compressed_tensors/transform/factory/random_hadamard.py
@@ -28,7 +28,14 @@ class RandomHadamardFactory(HadamardFactory):
     :param seed: random seed used to transform weight randomization
     """
 
-    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)
+    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)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
diff --git a/src/compressed_tensors/utils/helpers.py b/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,11 +374,23 @@ class ParameterizedDefaultDict(dict):
 
     def __init__(self, default_factory: Callable[[Any], Any]):
         self.default_factory = default_factory
+        self._factory_kwargs = frozendict()
 
-    def __missing__(self, key):
+    def __missing__(self, key: Any) -> Any:
         if isinstance(key, tuple):
-            value = self.default_factory(*key)
+            value = self.default_factory(*key, **self._factory_kwargs)
         else:
-            value = self.default_factory(key)
+            value = self.default_factory(key, **self._factory_kwargs)
         self[key] = value
         return value
+
+    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]

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ def _setup_packages() -> List:`
`88`	`88`	`)`
`89`	`89`
`90`	`90`	`def _setup_install_requires() -> List:`
`91`		`- return ["torch>=1.7.0", "transformers", "pydantic>=2.0"]`
	`91`	`+ return ["torch>=1.7.0", "transformers", "pydantic>=2.0", "frozendict"]`
`92`	`92`
`93`	`93`	`def _setup_extras() -> Dict:`
`94`	`94`	`return {`
Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,7 @@ def create_transform(self, module: Module, args: TransformArgs):`
`62`	`62`	`return RandomMatrixTransform(weight, args)`
`63`	`63`
`64`	`64`	`def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:`
	`65`	`+ # TODO: verify that weight is invertible (has non-zero determinant)`
`65`	`66`	`data = torch.rand(`
`66`	`67`	`(size, size), generator=self.generator, dtype=dtype, device=device`
`67`	`68`	`)`