| """Fixed orthogonal transforms for packed low-bit linear layers. |
| |
| The randomized Hadamard transform is block diagonal over the input feature |
| dimension. Both stored weight rows and row-major activations use the same |
| forward transform ``v -> (v * D) H``. Consequently, for orthogonal ``R = H D`` |
| |
| ``linear(x, W) == linear(transform(x), transform(W))``. |
| |
| Only the transformed weight is ternary and stored. The activation transform |
| must remain part of the deployed operator; materializing the inverse transform |
| into the weight would destroy the ternary representation. |
| """ |
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .quantization import hard_codes_scales |
|
|
|
|
| def _power_of_two(value: int) -> bool: |
| return value > 0 and value & (value - 1) == 0 |
|
|
|
|
| def normalized_hadamard(value: torch.Tensor) -> torch.Tensor: |
| """Apply an orthonormal Walsh-Hadamard transform on the last dimension.""" |
| features = int(value.shape[-1]) |
| if not _power_of_two(features): |
| raise ValueError("Hadamard dimension must be a positive power of two") |
| result = value |
| width = 1 |
| prefix = value.shape[:-1] |
| while width < features: |
| paired = result.reshape(*prefix, -1, 2, width) |
| left = paired[..., 0, :] |
| right = paired[..., 1, :] |
| result = torch.cat((left + right, left - right), dim=-1).reshape( |
| *prefix, features |
| ) |
| width *= 2 |
| return result / math.sqrt(features) |
|
|
|
|
| def rademacher_signs( |
| groups: int, |
| group_size: int, |
| *, |
| seed: int, |
| device: torch.device | str, |
| dtype: torch.dtype, |
| ) -> torch.Tensor: |
| """Generate device-independent deterministic Rademacher signs.""" |
| if groups < 1: |
| raise ValueError("groups must be positive") |
| if not _power_of_two(group_size): |
| raise ValueError("group_size must be a positive power of two") |
| generator = torch.Generator(device="cpu") |
| generator.manual_seed(int(seed)) |
| bits = torch.randint( |
| 0, |
| 2, |
| (groups, group_size), |
| generator=generator, |
| dtype=torch.int8, |
| device="cpu", |
| ) |
| return bits.to(device=device, dtype=dtype).mul(2).sub(1) |
|
|
|
|
| def blockwise_randomized_hadamard( |
| value: torch.Tensor, |
| *, |
| group_size: int = 128, |
| seed: int = 109, |
| ) -> torch.Tensor: |
| """Map row vectors to the randomized-Hadamard basis. |
| |
| The last dimension must be exactly divisible by ``group_size``. Padding is |
| deliberately rejected because padding before an orthogonal transform would |
| change the represented linear operator. |
| """ |
| features = int(value.shape[-1]) |
| if features % group_size: |
| raise ValueError("input features must be divisible by group_size") |
| groups = features // group_size |
| grouped = value.reshape(*value.shape[:-1], groups, group_size) |
| signs = rademacher_signs( |
| groups, |
| group_size, |
| seed=seed, |
| device=value.device, |
| dtype=value.dtype, |
| ) |
| return normalized_hadamard(grouped * signs).reshape_as(value) |
|
|
|
|
| def inverse_blockwise_randomized_hadamard( |
| value: torch.Tensor, |
| *, |
| group_size: int = 128, |
| seed: int = 109, |
| ) -> torch.Tensor: |
| """Map transformed weight rows back to the original dense basis.""" |
| features = int(value.shape[-1]) |
| if features % group_size: |
| raise ValueError("input features must be divisible by group_size") |
| groups = features // group_size |
| grouped = value.reshape(*value.shape[:-1], groups, group_size) |
| signs = rademacher_signs( |
| groups, |
| group_size, |
| seed=seed, |
| device=value.device, |
| dtype=value.dtype, |
| ) |
| return (normalized_hadamard(grouped) * signs).reshape_as(value) |
|
|
|
|
| class FixedTernaryLinear(nn.Module): |
| """Inference-only ternary linear with identity or randomized-Hadamard input.""" |
|
|
| def __init__( |
| self, |
| codes: torch.Tensor, |
| scales: torch.Tensor, |
| *, |
| compute_dtype: torch.dtype, |
| transform: str = "identity", |
| transform_seed: int = 109, |
| bias: torch.Tensor | None = None, |
| ): |
| super().__init__() |
| if codes.ndim != 3 or scales.shape != codes.shape[:2]: |
| raise ValueError("codes must be [out, groups, group_size]") |
| if not torch.all((codes >= -1) & (codes <= 1)): |
| raise ValueError("codes must be ternary") |
| if transform not in {"identity", "rht"}: |
| raise ValueError("transform must be 'identity' or 'rht'") |
| self.register_buffer("ternary_codes", codes.detach().to(torch.int8)) |
| self.register_buffer("group_scales", scales.detach().float()) |
| if bias is not None: |
| self.register_buffer("bias", bias.detach().clone()) |
| else: |
| self.bias = None |
| self.compute_dtype = compute_dtype |
| self.transform = transform |
| self.transform_seed = int(transform_seed) |
| self.in_features = int(codes.shape[1] * codes.shape[2]) |
| self.out_features = int(codes.shape[0]) |
| self.group_size = int(codes.shape[2]) |
| self.register_buffer( |
| "_evaluation_weight", |
| ( |
| self.ternary_codes.float() |
| * self.group_scales.unsqueeze(-1) |
| ) |
| .reshape(self.out_features, self.in_features) |
| .to(compute_dtype), |
| persistent=False, |
| ) |
|
|
| @classmethod |
| @torch.no_grad() |
| def from_weight( |
| cls, |
| weight: torch.Tensor, |
| *, |
| group_size: int = 128, |
| transform: str = "identity", |
| transform_seed: int = 109, |
| bias: torch.Tensor | None = None, |
| ) -> "FixedTernaryLinear": |
| if weight.ndim != 2: |
| raise ValueError("weight must be a matrix") |
| if weight.shape[1] % group_size: |
| raise ValueError("weight input features must be divisible by group_size") |
| transformed = weight.detach().float() |
| if transform == "rht": |
| transformed = blockwise_randomized_hadamard( |
| transformed, group_size=group_size, seed=transform_seed |
| ) |
| elif transform != "identity": |
| raise ValueError("transform must be 'identity' or 'rht'") |
| flat_codes, scales = hard_codes_scales(transformed, group_size) |
| codes = flat_codes.reshape(weight.shape[0], -1, group_size) |
| return cls( |
| codes, |
| scales, |
| compute_dtype=weight.dtype, |
| transform=transform, |
| transform_seed=transform_seed, |
| bias=bias, |
| ) |
|
|
| def transformed_weight(self) -> torch.Tensor: |
| return self._evaluation_weight |
|
|
| def effective_weight(self) -> torch.Tensor: |
| value = self.transformed_weight() |
| if self.transform == "rht": |
| value = inverse_blockwise_randomized_hadamard( |
| value, |
| group_size=self.group_size, |
| seed=self.transform_seed, |
| ) |
| return value.to(self.compute_dtype) |
|
|
| def forward(self, value: torch.Tensor) -> torch.Tensor: |
| if self.transform == "rht": |
| value = blockwise_randomized_hadamard( |
| value, |
| group_size=self.group_size, |
| seed=self.transform_seed, |
| ) |
| weight = self.transformed_weight().to(value.dtype) |
| bias = None if self.bias is None else self.bias.to(value.dtype) |
| return F.linear(value, weight, bias) |
|
|
| def code_histogram(self) -> dict[int, int]: |
| values, counts = torch.unique(self.ternary_codes.cpu(), return_counts=True) |
| result = {-1: 0, 0: 0, 1: 0} |
| result.update( |
| {int(value): int(count) for value, count in zip(values, counts)} |
| ) |
| return result |
|
|
|
|
| class TransformedProxyTernaryLinear(nn.Module): |
| """Training wrapper for hard-forward proxy codes in a transformed basis. |
| |
| ``matrix.effective_weight()`` is expected to return the deployed weight in |
| transform space. The input transform stays explicit so no dense inverse |
| weight is materialized during forward. |
| """ |
|
|
| def __init__( |
| self, |
| matrix: nn.Module, |
| *, |
| transform: str = "rht", |
| transform_seed: int = 109, |
| bias: torch.Tensor | None = None, |
| ): |
| super().__init__() |
| if transform not in {"identity", "rht"}: |
| raise ValueError("transform must be 'identity' or 'rht'") |
| if not hasattr(matrix, "effective_weight"): |
| raise TypeError("matrix must provide effective_weight()") |
| self.matrix = matrix |
| if bias is not None: |
| self.register_buffer("bias", bias.detach().clone()) |
| else: |
| self.bias = None |
| self.transform = transform |
| self.transform_seed = int(transform_seed) |
| self.in_features = int(matrix.in_features) |
| self.out_features = int(matrix.out_features) |
| self.group_size = int(matrix.group_size) |
|
|
| def forward(self, value: torch.Tensor) -> torch.Tensor: |
| if self.transform == "rht": |
| value = blockwise_randomized_hadamard( |
| value, |
| group_size=self.group_size, |
| seed=self.transform_seed, |
| ) |
| weight = self.matrix.effective_weight().to(value.dtype) |
| bias = None if self.bias is None else self.bias.to(value.dtype) |
| return F.linear(value, weight, bias) |
|
|
