remove fast_hadamard_transform dependency.

Signed-off-by: lkk12014402 <kaokao.lv@intel.com>

Author: lkk12014402

auto_round/experimental/transform/apply.py

@@ -52,19 +52,24 @@ def _apply_to_module(
 
     # create transform as submodule
     transform_name = "forward_hadamard"
-    transform = build_transform(**config.dict())
-    # module.register_module(transform_name, transform)
 
     if config.location == "input":
         from .triton.mxfp4 import mxfp4_forward_kernel_wrapper
 
+        transform = build_transform(
+            **config.dict(),
+            device="cpu",
+            precision=module.dtype,
+            location="input"
+        )
+
         def input_hook(_, args):
             input = args[0]
             # transform(input)
             orig_shape = input.shape
             x_flat = input.contiguous().flatten(end_dim=-2)
             qdq_input, _ = mxfp4_forward_kernel_wrapper(
-                x_flat, transform.get_transform_matrix(input.device, input.dtype)
+                x_flat, transform.weight
             )
             return qdq_input.reshape(orig_shape)
 
@@ -77,6 +82,12 @@ def input_hook(_, args):
         # fuse transform into weight
         assert hasattr(module, "weight")
 
+        transform = build_transform(
+            **config.dict(),
+            device=module.weight.device,
+            precision=module.weight.dtype,
+        )
+
         if config.need_calibration:
             # for training, the weight changes with every forward pass
             # for autoround tuning: patch wrapper linear qdq_weight func
@@ -93,7 +104,7 @@ def input_hook(_, args):
             # delattr(module, transform_name)
             # fuse transform into weight
             with torch.no_grad():
-                getattr(module, "weight").copy_(transform(module.weight.to("cuda")).to(module.weight.device))
+                getattr(module, "weight").copy_(transform(module.weight).to(module.weight.device))
 
     else:
         # TODO: apply transform to output/q/k

auto_round/experimental/transform/transforms.py

@@ -18,7 +18,9 @@
 
 import torch
 import torch.nn as nn
-from fast_hadamard_transform import hadamard_transform
+
+from auto_round.experimental.transform.utils.hadamard import deterministic_hadamard_matrix
+from auto_round.experimental.transform.utils.matrix import apply_transform_weight
 
 
 def filter_kwarg_dict(fn_or_method: Callable, kwarg_dict: Dict[str, Any]) -> Dict[str, Any]:
@@ -40,19 +42,43 @@ def remove_parametrizations(self) -> None:
 
 class HadamardTransform(nn.Module):
 
-    def __init__(self, transform_block_size: int = 32):
+    def __init__(
+        self,
+        transform_block_size: int = 32,
+        device: torch.device = None,
+        precision: torch.dtype = None,
+        location: str = "weight",
+        module_type: type[torch.nn.Module] = torch.nn.Linear,
+    ):
         super().__init__()
-        self.dim = transform_block_size
-        self.scale = 1 / math.sqrt(self.dim)
+        self.size = transform_block_size
+        self.scale = 1 / math.sqrt(self.size)
+        self.location = location
+        self.module_type = module_type
+        self.weight = self._create_weight(self.size, device, precision)
+
+    def _create_weight(
+        self,
+        size: int,
+        device: torch.device = None,
+        precision: torch.dtype = None,
+    ) -> torch.nn.Parameter:
+        data = deterministic_hadamard_matrix(size, precision, device) * self.scale
+        # TODO: implement SpinQuant, which rotation matrix is learnable
+        return nn.Parameter(data, requires_grad=False)
 
-    # @torch.no_grad()
     def forward(self, x: torch.Tensor):
         # Hadamard transform is it own inverse
-        x_shape = x.shape
-        return hadamard_transform(x.view(-1, self.dim), scale=self.scale).view(x_shape)
-
-    def get_transform_matrix(self, device: torch.device = None, dtype: torch.dtype = None):
-        return hadamard_transform(torch.eye(self.dim, device=device, dtype=dtype), scale=1 / math.sqrt(self.dim))
+        ori_shape = x.shape
+        x = x.view(-1, self.size)
+        return (
+            apply_transform_weight(
+                self.weight.to(device=x.device),
+                x.to(dtype=self.weight.dtype),
+                self.location,
+                self.module_type,
+            )
+        ).to(x.dtype).view(ori_shape)
 
 
 TRANSFORMS = {
@@ -64,3 +90,4 @@ def get_transform_matrix(self, device: torch.device = None, dtype: torch.dtype =
 def build_transform(transform_type: str, **transform_kwargs):
     transform = TRANSFORMS[transform_type]
     return transform(**filter_kwarg_dict(transform.__init__, transform_kwargs))
+

auto_round/experimental/transform/utils/__init__.py

@@ -0,0 +1,149 @@
+
+import math
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+
+REPO_PATH = Path(__file__).parent / "hadamards.safetensors"
+
+
+__all__ = ["random_hadamard_matrix", "deterministic_hadamard_matrix", "is_pow2"]
+
+
+# note that hadamard matrix multiplication reuses the code from
+# https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/transform/utils/hadamard.py
+
+
+def deterministic_hadamard_matrix(
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """
+    Construct an n-by-n Hadamard matrix, using Sylvester's construction.
+    `n` must be a power of 2.
+
+    Adapated from https://github.com/scipy/scipy/blob/v1.15.2/scipy/linalg/_special_matrices.py  # noqa: E501
+
+    :param size: order of the matrix, must be a power of 2
+    :param dtype: data type of matrix
+    :param device: device to construct matrix on
+    :return: hadamard matrix of size `size`
+    """
+    if size <= 0:
+        raise ValueError("Cannot construct deterministic hadamard of size <= 0")
+
+    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, device=device)
+
+    # Sylvester's construction
+    for _ in range(log2):
+        H = torch.vstack((torch.hstack((H, H)), torch.hstack((H, -H))))
+
+    return H
+
+
+def random_hadamard_matrix(
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = torch.device("cpu"),
+    gen: torch.Generator | None = None,
+) -> torch.Tensor:
+    """
+    Produces a randomly generated Hadamard matrix. Differs from
+    `deterministic_hadamard_matrix` in that this function supports non powers of 2
+    and randomization using a seeded generator
+
+    Adapated from https://github.com/facebookresearch/SpinQuant/blob/main/utils/hadamard_utils.py  # noqa: E501
+    Known matrices were retrieved from N. J. A. Sloane's Library of Hadamard Matrices http://www.neilsloane.com/hadamard/  # noqa: E501
+
+    :param size: The dimension of the hamadard matrix
+    :param dtype: data type of matrix
+    :param device: device to construct matrix on
+    :param gen: Optional generator random values
+    :return: randomly generated hadamard matrix
+    """
+    Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=dtype)  # cpu
+    Q = Q.to(device=device)
+    Q = Q * 2 - 1
+    Q = torch.diag(Q)
+    return _matmul_hadU(Q)
+
+
+def is_pow2(n: int) -> bool:
+    """
+    Check if a number is a power of 2
+
+    :param n: number to check
+    :return: True iff `n` is a power of 2
+    """
+    return n > 0 and (n & (n - 1) == 0)
+
+
+def _fetch_hadamard_divisor(
+    n: int,
+    dtype: torch.dtype,
+    device: torch.device = torch.device("cpu"),
+    file_path: str = REPO_PATH,
+) -> torch.Tensor | None:
+    """
+    Fetch a known hadamard matrix from the given file path. The returned matrix will
+    be of of size `k` such that `n / k` is a power of two. Return None if no such
+    matrix exists.
+
+    Note: This function reopens the safetensors file every time it is called.
+    This is technically inefficient, but a very small runtime cost and simpler
+    than forcing callers to manage the file open context
+
+    :param n: size of known hadamard matrix
+    :param dtype: data type to move fetched hadamard to
+    :param device: device to move fetched hadamard to
+    :return: a known hadamard matrix of size `n` if one exists, else None
+    """
+    open_device = torch.device("cpu") if device.type == "meta" else device
+    with safe_open(file_path, framework="pt", device=str(open_device)) as file:
+        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, device=device)
+
+    return None
+
+
+def _matmul_hadU(X: torch.Tensor) -> torch.Tensor:
+    size = X.size(0)
+    dtype = X.dtype
+    device = X.device
+
+    # Check if we have the determined hadamard matrix
+    hadK = _fetch_hadamard_divisor(size, dtype, device=device)
+    if hadK is None:
+        raise ValueError(f"Cannot construct random hadamard matrix of size {size}")
+    K = hadK.size(0)
+
+    # Reshape diag matrix with randomized -1/+1
+    input = X.clone().view(-1, size, 1)
+    output = input.clone()
+    while input.shape[1] > K:
+        input = input.view(input.shape[0], input.shape[1] // 2, 2, input.shape[2])
+        output = output.view(input.shape)
+        output[:, :, 0, :] = input[:, :, 0, :] + input[:, :, 1, :]
+        output[:, :, 1, :] = input[:, :, 0, :] - input[:, :, 1, :]
+        output = output.view(input.shape[0], input.shape[1], -1)
+        (input, output) = (output, input)
+    assert input.shape[1] == K
+    del output
+
+    # Do not explicitly repeat - OOM
+    # input = torch.bmm(
+    #     hadK.repeat(len(input), 1, 1).to(input.device).to(input.dtype), input)
+    # Use bcast instead
+    input = hadK.view(1, K, K).to(input) @ input
+
+    # normalize
+    return input.view(X.shape)

auto_round/experimental/transform/utils/hadamard.py

@@ -0,0 +1,95 @@
+import torch
+
+
+__all__ = ["apply_transform_weight"]
+
+# note that apply_transform_weight reuses some code from
+# https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/transform/utils/matrix.py
+
+def apply_transform_weight(
+    transform_weight: torch.Tensor,
+    value: torch.Tensor,
+    location: str,
+    module_type: type[torch.nn.Module],
+) -> torch.Tensor:
+    """
+    Using the transform location, apply the transform_weight to the
+    given value wrt linear weights. For more info on input and output transforms,
+    see `TransformLocation`
+
+    The following explains how weights should be applied to values according to location
+
+    let  x          be input activation
+         W          be weight,
+         yh, xh, Wh be transformed output, input, weight
+
+    note that
+         y  = (x W.T)        // torch.nn.Linear
+
+    Choose values for yh, xh, and Wh which incorporate matrix transforms
+
+    let  V, Vi      be transform matrices on input side
+         U, Ui      be transform matrices on output side
+
+    pick xh = (x V)
+         Wh = (U.T W Vi.T)
+         yh = (y U)
+
+    The following shows that `yh = (xh) (Wh).T` for the chosen values of yh, xh, and Wh
+
+    (xh) (Wh).T = (x V) (U.T W Vi.T).T
+                = (x V) (Vi W.T U)        // transpose matrix product identity
+                = (x W.T) U
+                = y U
+                = yh
+
+    :param transform_weight: transform weight to apply
+    :param value: value to apply transform_weight to
+    :param location: determines how weight should be applied
+    :param model_type: result of type(module), passed in to determine application of
+        weight transform
+    :return: value after transform_weight has been applied
+    """
+
+    if location == "input":
+        return _multihead_matmul(value, transform_weight)
+
+    if module_type == torch.nn.Linear:
+        return _multihead_matmul(value, transform_weight.T)
+
+
+def _multihead_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    """
+    Performs A @ B for last two dims of two matrices A and B that possibly
+    have different shapes, as is the case in multi-headed dimension. If
+    shapes are different, this is equivalent to converting the last two dims
+    of the smaller matrix into a block-diagonal matrix with the same shape as
+    the last two dims of the larger matrix.
+
+    E.g. if A is half the size of B, this function will perform
+    [[A  ]  @ B
+     [  A]]
+
+    If B is a third of the size of A, this function will perform
+    A @ [[B    ]
+         [  B  ]
+         [    B]]
+
+    This function will error out if the shapes are not evenly divisble
+
+    :param A: left-hand tensor
+    :param B: right-hand tensor
+    :return: result
+    """
+    if A.shape[-1] > B.shape[-2]:
+        head_dim = B.shape[-2]
+        num_heads = A.shape[-1] // head_dim
+        A = A.unflatten(-1, (num_heads, head_dim))
+        return (A @ B).flatten(-2, -1)
+    elif A.shape[-1] < B.shape[-2]:
+        head_dim = A.shape[-1]
+        num_heads = B.shape[-2] // head_dim
+        B = B.unflatten(-2, (num_heads, head_dim))
+        return (A @ B).flatten(-3, -2)
+    else:
+        return A @ B