fix(rotation): persist rotation config and apply online Hadamard for inference and calibration (#2980)

* fix(rotation): persist rotation config and apply online Hadamard for inference

Cherry-pick/adapt from GPT-QModel-Ultra#66:
- Persist  in  and round-trip via .
- Fuse exact Hadamard into  weights during rotation.
- Add JIT  Hadamard extension in  (vendored
  from Dao-AILab/fast-hadamard-transform, BSD-3-Clause) plus .
- Attach // to  QuantLinear modules at
   and after . Restrict rotation to /.
- Call  in  and .
- Disable  when  is enabled.
- Add  and README/MANIFEST attribution/license.

* fix(rotation): wire online Hadamard into loader, base, qlinear, config, version, and attributions

* fix(rotation): apply online Hadamard during GPTQ calibration for mlp.down_proj

- Extract apply_online_hadamard so HookedLinear and BaseQuantLinear share the same transform.
- rotate_mlp_output now sets online_full_had/had_K/K on the nn.Linear before it becomes HookedLinear.
- HookedLinear.from_linear copies rotation state and forward applies H to the activation used for GPTQ calibration.
- Add unit test verifying HookedLinear cancellation and hook input transform.

Author: Qubitium

MANIFEST.in

@@ -6,6 +6,7 @@ recursive-include gptqmodel_ext/machete *.h *.hpp *.cuh *.cu *.cpp *.py
 recursive-include gptqmodel_ext/cutlass_extensions *.h *.hpp *.cuh *.cu *.cpp *.py
 recursive-include gptqmodel_ext/paroquant *.h *.cuh *.cu *.cpp
 recursive-include gptqmodel_ext/qqq *.h *.cuh *.cu *.cpp
+recursive-include gptqmodel_ext/hadamard *.h *.cuh *.cu *.cpp LICENSE
 recursive-include gptqmodel/exllamav3/util/hadamard_data *.txt
 include licenses/*
 include gptqmodel_ext/pack_block_cpu.cpp

README.md

@@ -666,6 +666,7 @@ Models quantized by GPT-QModel are inference compatible with HF Transformers (mi
 * EoRA: Nvidia, main-author: Shih-Yang Liu, arXiv preprint arXiv:2410.21271.
 * GAR: Intel, main-author: T Gafni, A Karnieli, Y Hanani, [Paper](https://openaccess.thecvf.com/content/CVPR2025W/eLVM/html/Gafni_Dual_Precision_Quantization_for_Efficient_and_Accurate_Deep_Neural_Networks_CVPRW_2025_paper.html)
 * GPTAQ: Yale Intelligent Computing Lab, main-author: Yuhang Li, arXiv:2504.02692.
+* Fast Hadamard Transform: [Dao-AILab/fast-hadamard-transform](https://github.com/Dao-AILab/fast-hadamard-transform), by Tri Dao, vendored into `gptqmodel_ext/hadamard/` under the BSD-3-Clause license. A copy of the license is included in `gptqmodel_ext/hadamard/LICENSE`.
 * QQQ: Meituan, main-author Ying Zhang, arXiv:2406.09904
 * FOEM: Zheng, Xingyu and Qin, Haotong and Li, Yuye and Chu, Haoran and Wang, Jiakai and Guo, Jinyang and Magno, Michele and Liu, Xianglong [Paper](https://ojs.aaai.org/index.php/AAAI/article/view/40123)
 

gptqmodel/extension.py

@@ -100,6 +100,13 @@ def _resolve_extension_factory(module_name: str, attr_name: str) -> TorchOpsJitE
         aliases=("paroquant_rotation",),
         resolve=lambda: _resolve_extension_attr("gptqmodel.utils.paroquant", "_PAROQUANT_ROTATION_EXTENSION"),
     ),
+    _ExtensionSpec(
+        name="hadamard",
+        aliases=("fast_hadamard_transform",),
+        resolve=lambda: _resolve_extension_attr("gptqmodel.utils.hadamard", "_HADAMARD_TORCH_OPS_EXTENSION"),
+        supported=lambda: _resolve_attr("gptqmodel.utils.hadamard", "hadamard_supported")(),
+        unsupported_error=lambda: _resolve_attr("gptqmodel.utils.hadamard", "hadamard_runtime_error")(),
+    ),
 )
 
 _EXTENSION_SPECS_BY_NAME = {spec.name: spec for spec in _EXTENSION_SPECS}

gptqmodel/models/base.py

@@ -94,7 +94,7 @@
     EXPERT_INDEX_PLACEHOLDER,
     META,
 )
-from .loader import ModelLoader
+from .loader import ModelLoader, _setup_rotation_online_had
 from .writer import ModelWriter
 
 
@@ -1001,6 +1001,15 @@ def quantize(
             self.model, _ = rotate_model(model=self.model, rotate_mode=self.quantize_config.rotation,
                                             device=rotation_device, **module_name_args)
 
+        if self.quantize_config.rotation:
+            backend = normalize_backend(backend, quant_method=self.quantize_config.method)
+            if backend == BACKEND.AUTO:
+                backend = BACKEND.GPTQ_TORCH
+            if backend not in (BACKEND.GPTQ_TORCH, BACKEND.GPTQ_TRITON):
+                raise NotImplementedError(
+                    f"`rotation` is only supported with `gptq_torch` or `gptq_triton` backend, got `{backend}`."
+                )
+
         if self.quantize_config.uses_weight_only_lifecycle():
             result = self._quantize_weight_only(
                 calibration=calibration,
@@ -1029,6 +1038,7 @@ def quantize(
         if timer is not None:
             timer.flush()
 
+        _setup_rotation_online_had(self.model, self.quantize_config.rotation)
         return result
 
     def _quantize_with_calibration(

gptqmodel/models/loader.py

@@ -216,6 +216,46 @@ def _set_paged_attention_safe_cuda_graphs(model) -> None:
         cb_config.use_cuda_graph = (False, False)
 
 
+def _setup_rotation_online_had(model, rotation: Optional[str]) -> None:
+    """Attach online Hadamard transform state to QuantLinear modules for rotation inference.
+
+    Rotation (QuaRot/SpinQuant/Hadamard) fuses an orthogonal ``Q`` and an exact
+    Hadamard ``H`` into the weights. The ``H`` must be re-applied to activations
+    at inference before the quantized matmul for ``mlp.down_proj``. The per-head
+    Hadamard fused into ``v_proj``/``o_proj`` cancels inside the attention path,
+    so ``self_attn.o_proj`` does not require an online transform. Only the
+    ``rotation`` string is persisted; ``had_K``/``K`` are recomputed from the
+    model dimensions at load time.
+    """
+    if not rotation:
+        return
+
+    from ..quantization.rotation.hadamard_utils import get_hadK
+
+    if rotation not in {"hadamard", "random"}:
+        raise ValueError(f"Unsupported rotation mode: `{rotation}`")
+
+    online_count = 0
+    for name, module in model.named_modules():
+        if not isinstance(module, BaseQuantLinear):
+            continue
+
+        if name.endswith("mlp.down_proj"):
+            intermediate_size = module.in_features
+            had_K, K = get_hadK(intermediate_size)
+            module.online_full_had = True
+            module.K = K
+            if had_K is not None:
+                module.register_buffer("had_K", had_K, persistent=False)
+            online_count += 1
+
+    if online_count == 0:
+        log.warn(
+            f"Rotation `{rotation}` requested but no `mlp.down_proj` QuantLinear "
+            "modules were found; online Hadamard transform was not applied."
+        )
+
+
 def _is_accelerated_attention_device(device: object) -> bool:
     """Return True when the selected device can run CUDA/ROCm flash attention."""
 
@@ -1302,6 +1342,19 @@ def from_quantized(
                     log.info(f"The layer {name} is not quantized.")
                     del modules[name]
 
+            if qcfg.rotation:
+                if format_code not in (FORMAT.GPTQ, FORMAT.GPTQ_V2):
+                    raise NotImplementedError(
+                        f"`rotation` is only supported for GPTQ/GPTQ_V2 checkpoints, got `{format_code}`."
+                    )
+                backend = normalize_backend(backend, quant_method=qcfg.method)
+                if backend == BACKEND.AUTO:
+                    backend = BACKEND.GPTQ_TORCH
+                if backend not in (BACKEND.GPTQ_TORCH, BACKEND.GPTQ_TRITON):
+                    raise NotImplementedError(
+                        f"`rotation` is only supported with `gptq_torch` or `gptq_triton` backend, got `{backend}`."
+                    )
+
             if format_code == FORMAT.EXL3:
                 if not isinstance(qcfg.tensor_storage, dict) or not qcfg.tensor_storage:
                     raise ValueError("EXL3 checkpoints require `quantization_config.tensor_storage` metadata.")
@@ -1728,6 +1781,7 @@ def assign(mod, device_id):
             trust_remote_code=trust_remote_code,
             model_local_path=model_local_path,
         )
+        _setup_rotation_online_had(instance.model, qcfg.rotation)
         _set_paged_attention_safe_cuda_graphs(instance.model)
         return instance
 

gptqmodel/nn_modules/hooked_linear.py

@@ -10,6 +10,7 @@
 from accelerate.utils import has_offloaded_params
 from torch import nn
 
+from ..quantization.rotation.hadamard_utils import apply_online_hadamard
 from ..utils.device_telemetry import emit_device_telemetry
 from ..utils.logger import setup_logger
 
@@ -248,6 +249,11 @@ def from_linear(linear: torch.nn.Linear):
         custom_linear = HookedLinear(linear.in_features, linear.out_features)
         custom_linear.weight = linear.weight
         custom_linear.bias = linear.bias
+        custom_linear.online_full_had = getattr(linear, "online_full_had", False)
+        custom_linear.online_partial_had = getattr(linear, "online_partial_had", False)
+        custom_linear.had_dim = getattr(linear, "had_dim", -1)
+        custom_linear.had_K = getattr(linear, "had_K", None)
+        custom_linear.K = getattr(linear, "K", 1)
         return custom_linear
 
     @torch.inference_mode()
@@ -264,6 +270,14 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         )
         if original_device != target_device:
             input = input.to(device=target_device)
+        input = apply_online_hadamard(
+            input,
+            online_full_had=getattr(self, "online_full_had", False),
+            online_partial_had=getattr(self, "online_partial_had", False),
+            had_K=getattr(self, "had_K", None),
+            K=getattr(self, "K", 1),
+            had_dim=getattr(self, "had_dim", -1),
+        )
         output = super().forward(input)
         if self.forward_hook:
             self.forward_hook(self, (input,), output)

gptqmodel/nn_modules/qlinear/__init__.py

@@ -121,6 +121,13 @@ def __init__(self,
         self._autotune_complete = False
         self._autotune_result: Any = None
 
+        # Rotation/online-Hadamard state for SpinQuant/QuaRot inference.
+        self.online_full_had = False
+        self.online_partial_had = False
+        self.had_dim = -1
+        self.K = 1
+        self.register_buffer("had_K", None, persistent=False)
+
         validate_args = {
             "bits": bits,
             "in_features": in_features,
@@ -463,6 +470,19 @@ def train(self, mode=True):
         self.clear_autotune()
         return super().train(mode)
 
+    def _apply_rotation_to_input(self, x: t.Tensor) -> t.Tensor:
+        """Apply online Hadamard transform to the input for SpinQuant/QuaRot inference."""
+        from ...quantization.rotation.hadamard_utils import apply_online_hadamard
+
+        return apply_online_hadamard(
+            x,
+            online_full_had=self.online_full_had,
+            online_partial_had=self.online_partial_had,
+            had_K=self.had_K,
+            K=self.K,
+            had_dim=self.had_dim,
+        )
+
 
 class GroupedQuantLinear(BaseQuantLinear):
     SUPPORTS_GROUP_SIZE: List[int] = None

gptqmodel/nn_modules/qlinear/torch.py

@@ -306,6 +306,7 @@ def forward(self, x: torch.Tensor):
         return out
 
     def _forward(self, x, out_shape):
+        x = self._apply_rotation_to_input(x)
         cached = self._maybe_get_cached_weights(x)
         if cached is not None:
             out = torch.matmul(x, cached).reshape(out_shape)

gptqmodel/nn_modules/qlinear/tritonv2.py

@@ -179,6 +179,8 @@ def forward(self, x):
         if self.training:
             return super().forward(x)
 
+        x = self._apply_rotation_to_input(x)
+
         # if in_features is padded, we need to pad the input as well
         # if x.size(-1) != self.padded_infeatures:
         #     x = F.pad(x, (0, self.padded_infeatures - self.in_features))

gptqmodel/quantization/config.py

@@ -2556,6 +2556,13 @@ def __post_init__(self):
             self.meta = {}
 
         self.adapter = normalize_adapter(self.adapter)
+
+        # Rotation fuses orthogonal transforms into the weights and requires
+        # materialized tensors; meta-device/shell loading cannot be used.
+        if self.rotation and self.offload_to_disk:
+            log.warn(f"{self.__class__.__name__}: `rotation` is incompatible with `offload_to_disk`; disabling disk offload.")
+            self.offload_to_disk = False
+
         self._ensure_offload_temp_dir()
 
         self.dense_vram_strategy = _normalize_dense_vram_strategy(self.dense_vram_strategy)
@@ -2940,6 +2947,7 @@ def to_dict(self):
             FORMAT_FIELD_CODE: self.format,
             FORMAT_FIELD_CHECKPOINT: self.format,
             PACK_DTYPE_FIELD: str(self.pack_dtype).split(".")[-1],
+            "rotation": self.rotation,
             META_FIELD: meta_payload,
         }
         self._update_output_payload(out)
@@ -3732,6 +3740,13 @@ def __post_init__(self):
             self.meta = {}
 
         self.adapter = normalize_adapter(self.adapter)
+
+        # Rotation fuses orthogonal transforms into the weights and requires
+        # materialized tensors; meta-device/shell loading cannot be used.
+        if self.rotation and self.offload_to_disk:
+            log.warn(f"{self.__class__.__name__}: `rotation` is incompatible with `offload_to_disk`; disabling disk offload.")
+            self.offload_to_disk = False
+
         self._ensure_offload_temp_dir()
 
         self.dense_vram_strategy = _normalize_dense_vram_strategy(self.dense_vram_strategy)