solidify dtype, add gpu tests

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

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -59,7 +59,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         return HadamardTransform(weight, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
-        data = deterministic_hadamard_matrix(size)
+        data = deterministic_hadamard_matrix(size, dtype=dtype)
         data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 

src/compressed_tensors/transform/factory/random_hadamard.py

@@ -29,6 +29,6 @@ class RandomHadamardFactory(HadamardFactory):
     """
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
-        data = random_hadamard_matrix(size, self.generator)
+        data = random_hadamard_matrix(size, dtype=dtype, gen=self.generator)
         data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)

src/compressed_tensors/transform/utils/hadamard.py

@@ -21,7 +21,6 @@
 
 
 REPO_PATH = os.path.join(os.path.dirname(__file__), "hadamards.safetensors")
-DTYPE = torch.int32
 
 
 __all__ = ["random_hadamard_matrix", "deterministic_hadamard_matrix", "is_pow2"]
@@ -31,7 +30,9 @@
 # https://github.com/Dao-AILab/fast-hadamard-transform/tree/master
 
 
-def deterministic_hadamard_matrix(size: int) -> torch.Tensor:
+def deterministic_hadamard_matrix(
+    size: int, dtype: torch.dtype = torch.bfloat16
+) -> torch.Tensor:
     """
     Construct an n-by-n Hadamard matrix, using Sylvester's construction.
     `n` must be a power of 2.
@@ -44,11 +45,11 @@ def deterministic_hadamard_matrix(size: int) -> torch.Tensor:
     if size <= 0:
         raise ValueError("Cannot construct deterministic hadamard of size <= 0")
 
-    log2 = int(math.log(size, 2))
+    log2 = int(math.log2(size))
     if size != 2**log2:
         raise ValueError("Cannot construct deterministic hadamard of size != 2^n")
 
-    H = torch.tensor([[1]], dtype=DTYPE)
+    H = torch.tensor([[1]], dtype=dtype)
 
     # Sylvester's construction
     for _ in range(0, log2):
@@ -58,7 +59,9 @@ def deterministic_hadamard_matrix(size: int) -> torch.Tensor:
 
 
 def random_hadamard_matrix(
-    size: int, gen: Optional[torch.Generator] = None
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    gen: Optional[torch.Generator] = None,
 ) -> torch.Tensor:
     """
     Produces a randomly generated Hadamard matrix.
@@ -72,7 +75,7 @@ def random_hadamard_matrix(
     :return: randomly generated hadamard matrix
     """
     # Benefits: support other shapes / non powers of 2, support randomization
-    Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=DTYPE)
+    Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=dtype)
     Q = Q * 2 - 1
     Q = torch.diag(Q)
     return _matmul_hadU(Q) / math.sqrt(size)
@@ -82,7 +85,9 @@ def is_pow2(n: int) -> bool:
     return (n & (n - 1) == 0) and (n > 0)
 
 
-def _get_known_divisor(n: int, file_path: str = REPO_PATH) -> Optional[torch.Tensor]:
+def _get_known_divisor(
+    n: int, dtype: torch.dtype, file_path: str = REPO_PATH
+) -> Optional[torch.Tensor]:
     """
     Fetch a known hadamard matrix from the given file path. The returned matrix will
     be of of size `k` such that `n` divides `d` and `n / d` is a power of two. Return
@@ -100,16 +105,17 @@ def _get_known_divisor(n: int, file_path: str = REPO_PATH) -> Optional[torch.Ten
         divisors = sorted([int(key) for key in file.keys()], reverse=True)
         for divisor in divisors:
             if n % divisor == 0 and is_pow2(n // divisor):
-                return file.get_tensor(str(divisor)).to(dtype=DTYPE)
+                return file.get_tensor(str(divisor)).to(dtype=dtype)
 
     return None
 
 
 def _matmul_hadU(X: torch.Tensor) -> torch.Tensor:
     size = X.shape[-1]
+    dtype = X.dtype
 
     # Check if we have the determined hadamard matrix
-    hadK = _get_known_divisor(size)
+    hadK = _get_known_divisor(size, dtype)
     if hadK is None:
         raise ValueError(f"Cannot construct random hadamard matrix of size {size}")
     K = hadK.size(0)

tests/test_transform/utils/test_hadamard.py

@@ -20,6 +20,7 @@
     is_pow2,
     random_hadamard_matrix,
 )
+from tests.testing_utils import requires_gpu
 
 
 _sizes_to_test = [
@@ -28,27 +29,29 @@
     1280,  # qwen_2_5_vl vision
     1600,  # gpt2 xl
     2048,  # gpt3 small
-    # 3584,  # qwen_2_5_vl
-    # 3840,  # qwen_2_5_vl vision qkv
-    # 4096,  # llama3
-    # 14336,  # llama3 intermediate
-    # 18944,  # qwen_2_5_vl intermediate
+    3584,  # qwen_2_5_vl
+    3840,  # qwen_2_5_vl vision qkv
+    4096,  # llama3
+    14336,  # llama3 intermediate
+    18944,  # qwen_2_5_vl intermediate
 ]
 
 
+@requires_gpu
 @pytest.mark.parametrize("size", _sizes_to_test)
 def test_random_hadamard_matrix_compliant(size):
     # (H / sqrt(n))(H.T / sqrt(n)) == I
-    had_matrix = random_hadamard_matrix(size)
-    product = torch.round(had_matrix @ had_matrix.T)
-    assert torch.allclose(product, torch.eye(size, dtype=product.dtype), atol=1e-5)
+    with torch.device("cuda"):
+        had_matrix = random_hadamard_matrix(size)
+        product = torch.round(had_matrix @ had_matrix.T)
+        assert torch.allclose(product, torch.eye(size, dtype=product.dtype), atol=1e-5)
 
 
 def test_random_hadamard_generator():
     # check that generation is deterministic with a seed
     generator = torch.Generator().manual_seed(42)
-    one = random_hadamard_matrix(2048, generator)
-    two = random_hadamard_matrix(2048, generator)
+    one = random_hadamard_matrix(2048, gen=generator)
+    two = random_hadamard_matrix(2048, gen=generator)
 
     one_true = torch.tensor(
         [
@@ -69,14 +72,16 @@ def test_random_hadamard_generator():
     assert torch.all(two[:3, :3].sign() == two_true.sign())
 
 
+@requires_gpu
 @pytest.mark.parametrize("size", _sizes_to_test)
 def test_deterministic_hadamard_compliant(size):
-    if not is_pow2(size):
-        with pytest.raises(ValueError):
-            had_matrix = deterministic_hadamard_matrix(size)
-        return
+    with torch.device("cuda"):
+        if not is_pow2(size):
+            with pytest.raises(ValueError):
+                had_matrix = deterministic_hadamard_matrix(size)
+            return
 
-    # (H / sqrt(n))(H.T / sqrt(n)) == I
-    had_matrix = deterministic_hadamard_matrix(size)
-    product = had_matrix @ had_matrix.T
-    assert torch.allclose(product, torch.eye(size, dtype=product.dtype), atol=1e-5)
+        # (H / sqrt(n))(H.T / sqrt(n)) == I
+        had_matrix = deterministic_hadamard_matrix(size)
+        product = had_matrix @ had_matrix.T
+        assert torch.allclose(product, torch.eye(size, dtype=product.dtype), atol=1e-5)