Support 3D Weights in GPTQ Algorithm (#3835)

### Changes

The approach here is quite straightfoward, `_quantize_weights` works as
usual for 2D weights. The difference is in calculate hessian where the
hessian is 3D in both 3D and 2D weights case. By default hessian has the
shape (1, hidden_dim, hidden_dim).
Before this was just (hidden_dim, hidden_dim). For 3D, it is
(num_experts/batch, hidden_dim, hidden_dim).

Now, this 3D hessian or "batched" hessian is looped over and the 2D
weight is extracted and passed to the old `_quantize_weights` function
as usual and scale/zp are returned. These scales and zp are then stacked
together in a collector variable. For 2D case, it is flattened. For 3D
the stacked scale, zp are returned.

NOTE: Scale Estimation + GPTQ support is not added for 3D weights yet

### Reason for changes

Support 3D weights for models like MoE in GPTQ

### Related tickets

175789 & 175212

### Tests

Model: Qwen/Qwen3-30B-A3B
NNCF Backend: OpenVINO
Higher is better.
Task: gsm8k
Limit: 100
Max New Tokens: 10000
OpenVINO version: 2026.0.0.dev20251111 (with WA for 176465)
n-shots: 5(default)

Precision Type | Filter | Value
-- | -- | --
INT4 SYM GS128 (with GPTQ) Calibrated on GSM8k with 128 samples |
flexible-extract | 0.79
  | strict-match | 0.64
INT4 SYM GS128 (with GPTQ after bug fix commit
dc355fe)
Calibrated on GSM8k with 128 samples | flexible-extract | 0.78
  | strict-match | 0.57
INT4 SYM GS128 | flexible-extract | 0.55 
  | strict-match | 0.29 
FP32 | flexible-extract | 0.92 
  | strict-match | 0.82

---------

Co-authored-by: Lyalyushkin Nikolay <nikolay.lyalyushkin@intel.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

src/nncf/quantization/algorithms/weight_compression/gptq.py

@@ -10,6 +10,8 @@
 # limitations under the License.
 
 import math
+from functools import reduce
+from operator import mul
 from typing import Optional, TypeVar
 
 import nncf
@@ -130,8 +132,39 @@ def apply(
                 raise nncf.UnsupportedModelError(msg)
 
             _, input_tensors = next(iter(inputs.items()))
-            hessian = self._calculate_hessian(node, input_tensors)
-            scale, zero_point = self._quantize_weights(model, graph, wc_params, hessian, input_tensors)
+            weight_tensor = self._backend_entity.get_weight(
+                wc_params.node_with_weight, wc_params.weight_port_id, model, graph
+            )
+            weight_tensor = fns.astype(weight_tensor, TensorDataType.float32)
+
+            is_3d_weight = len(weight_tensor.shape) == 3
+
+            node = wc_params.node_with_weight
+            hessian = self._calculate_hessian(node, input_tensors, is_3d_weight)
+            weight_tensor = fns.unsqueeze(weight_tensor, 0) if not is_3d_weight else weight_tensor
+            scales = []
+            zero_points = []
+            weights = []
+            for batch_idx in range(hessian.shape[0]):
+                batch_hessian = hessian[batch_idx]
+                batch_weight = weight_tensor[batch_idx]
+                reduction_axes = wc_params.reduction_axes
+                assert len(reduction_axes) == 1, "2D reduction axes is not currently supported in GPTQ"
+                wc_params.reduction_axes = (reduction_axes[0] - 1,) if is_3d_weight else reduction_axes
+                # Input tensors is a List of tensors with shape [batch_size, seq_len, hidden_dim] for 3D weights case
+                # So we need to prepare the list by selecting only the current batch inputs only
+                input_tensor = [inp[batch_idx] for inp in input_tensors] if is_3d_weight else input_tensors
+                batch_quantized_weight, batch_scale, batch_zero_point = self._quantize_weights(
+                    wc_params, batch_hessian, batch_weight, input_tensor
+                )
+                wc_params.reduction_axes = reduction_axes
+                weights.append(batch_quantized_weight)
+                scales.append(batch_scale)
+                zero_points.append(batch_zero_point)
+            scale = fns.stack(scales, axis=0) if is_3d_weight else scales[0]
+            zero_point = fns.stack(zero_points, axis=0) if is_3d_weight and None not in zero_points else zero_points[0]
+            weight = fns.stack(weights, axis=0) if is_3d_weight else weights[0]
+            self._backend_entity.set_weight(wc_params.node_with_weight, wc_params.weight_port_id, model, graph, weight)
             res[wc_params.weight_name] = CompressedWeight(None, scale, zero_point, None)
 
         return model, res
@@ -163,7 +196,7 @@ def get_statistic_points(
 
         return self._layerwise_engine.get_statistic_points(model, graph, filtered_nodes)
 
-    def _calculate_hessian(self, node: NNCFNode, inputs: list[Tensor]) -> Tensor:
+    def _calculate_hessian(self, node: NNCFNode, inputs: list[Tensor], is_3d_weight: bool = False) -> Tensor:
         """
         Calculates the Hessian matrix for the given node and inputs.
 
@@ -179,30 +212,39 @@ def _calculate_hessian(self, node: NNCFNode, inputs: list[Tensor]) -> Tensor:
         if node.layer_attributes.input_attributes["transpose"]:
             msg = "Transposed input is not supported"
             raise nncf.UnsupportedModelError(msg)
-
+        # Make hessian 3D. Such that for 2D weights it is only 1 batch and can be squeezed later.
+        # For 3D weights this dimension matches the weights dimensions
+        hessian_batch = 1 if not is_3d_weight else reduce(mul, inputs[0].shape[:-2])
         hessian = fns.zeros(
-            (inputs[0].shape[-1], inputs[0].shape[-1]), backend=inputs[0].backend, dtype=TensorDataType.float32
+            (hessian_batch, inputs[0].shape[-1], inputs[0].shape[-1]),
+            backend=inputs[0].backend,
+            dtype=TensorDataType.float32,
         )
 
         for inp in inputs:
-            batch_size = 1 if len(inp.shape) == 2 else inp.shape[0]
+            is_3d_act = len(inp.shape) == 3
+            # For 3D weights case, batch size will always be 1. Each "batch"/expert of the activation is treated as
+            # single 2D matmuls
+            batch_size = 1 if is_3d_weight or not is_3d_act else inp.shape[0]
             if node.metatype in self._backend_entity.matmul_metatypes:
-                if len(inp.shape) == 3:
+                # For 3D act + 2D weight case we should reshape activation to 2D to match weight
+                # For 3D act + 3D weight it should remain in 3D and the last 2 dimensions should be activation per
+                # batch/0-th dimension
+                if is_3d_act and not is_3d_weight:
                     inp = inp.reshape((-1, inp.shape[-1]))
-                inp = fns.transpose(inp)
+                inp = fns.moveaxis(inp, -1, -2)
             hessian *= nsamples / (nsamples + batch_size)
             nsamples += batch_size
             inp = fns.astype(inp, TensorDataType.float32) * math.sqrt(2 / nsamples)
-            hessian += fns.matmul(inp, fns.transpose(inp))
+            hessian += fns.matmul(inp, fns.moveaxis(inp, -1, -2))
 
         return hessian
 
     def _quantize_weights(
         self,
-        model: TModel,
-        graph: NNCFGraph,
         wc_params: WeightCompressionParameters,
         hessian: Tensor,
+        weight_tensor: Tensor,
         inputs: list[Tensor],
     ):
         """
@@ -221,10 +263,7 @@ def _quantize_weights(
             msg = "Transpose is not supported"
             raise RuntimeError(msg)
 
-        weight_tensor = self._backend_entity.get_weight(
-            wc_params.node_with_weight, wc_params.weight_port_id, model, graph
-        )
-        weight_tensor = fns.astype(weight_tensor, TensorDataType.float32)
+        assert len(hessian.shape) == 2, "Hessian should be 2D"
 
         dead_indices = fns.diag(hessian) == 0
         hessian[dead_indices, dead_indices] = 1
@@ -323,9 +362,6 @@ def _quantize_weights(
             weight_tensor[:, i2:] -= fns.matmul(error_block, hessian_inv[i1:i2, i2:])
 
         quantized_tensor = quantized_tensor.reshape(weight_tensor.shape).astype(weight_tensor.dtype)
-        self._backend_entity.set_weight(
-            wc_params.node_with_weight, wc_params.weight_port_id, model, graph, quantized_tensor
-        )
 
         scales = fns.stack(scales, axis=1)
         if wc_params.compression_config.group_size == -1:
@@ -339,4 +375,4 @@ def _quantize_weights(
                 zero_points = fns.squeeze(zero_points, axis=-1)
         else:
             zero_points = None
-        return scales, zero_points
+        return quantized_tensor, scales, zero_points

tests/openvino/native/quantization/test_gptq.py

@@ -16,8 +16,10 @@
 
 import numpy as np
 import openvino as ov
+import pytest
 import torch
 
+from nncf import Dataset
 from nncf.common.factory import build_graph
 from nncf.parameters import CompressWeightsMode
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionConfig
@@ -323,53 +325,128 @@ def fasterquant(
         return scale, zero, g_idx
 
 
-def test_calculate_scale_linear():
-    # generate inputs
+class Linear3DModel(torch.nn.Module):
+    def __init__(self, weight_3d: np.ndarray):
+        super().__init__()
+        w = torch.from_numpy(weight_3d)
+        self.weight = torch.nn.Parameter(w)
+
+    def forward(self, x):
+        # OV expects transposed constant when applying GPTQ
+        return torch.matmul(x, self.weight.transpose(1, 2))
+
+
+def _create_ov_model(weights: np.ndarray, input_shape: tuple, is_3d_weights: bool = False):
+    import openvino.runtime.opset13 as opset
+
+    param = opset.parameter(input_shape, dtype=np.float32, name="input")
+    const = opset.constant(weights, dtype=np.float32, name="self.weight")
+    matmul = opset.matmul(param, const, transpose_a=False, transpose_b=True)
+    result = opset.result(matmul, name="output")
+    return ov.Model([result], [param])
+
+
+def _make_nncf_dataset(ov_model, inputs: list[np.ndarray]) -> Dataset:
+    input_name = ov_model.inputs[0].get_any_name()
+    items = [{input_name: inp} for inp in inputs]
+    return Dataset(items, lambda x: x)
+
+
+@pytest.mark.parametrize("is_3d_weights", [False, True], ids=["2D_weights", "3D_weights"])
+def test_calculate_scale_linear(is_3d_weights: bool):
     np.random.seed(0)
-    inputs = [np.random.rand(128, 32).astype(np.float32) for _ in range(10)]
-    weights = np.random.rand(20, 32).astype(np.float32)
 
-    # calculate reference
-    with torch.no_grad():
-        layer = torch.nn.Linear(32, 20)
-        layer.weight.copy_(torch.from_numpy(weights))
+    hidden_dims = 32
+    out_dims = 20
+    group_size = 16
+    n_inputs = 10
+    batch_size = 4
 
-        ref_gptq = GPTQReference(layer)
-        for inp in inputs:
-            ref_gptq.add_batch(torch.from_numpy(inp))
+    inputs = [np.random.rand(batch_size, 128, hidden_dims).astype(np.float32) for _ in range(n_inputs)]
+    weights = np.random.rand(batch_size, out_dims, hidden_dims).astype(np.float32)
 
-        ref_scale, _, _ = ref_gptq.fasterquant(percdamp=0.1, group_size=16)
+    # Select only first batch for 2D case and make it 2D
+    inputs = inputs if is_3d_weights else [activation[0] for activation in inputs]
+    weights = weights if is_3d_weights else weights[0]
 
-    # convert PyTorch model to OpenVINO
-    ov_model = ov.convert_model(layer, example_input=inputs[0])
+    # Step 1: We apply a reference GPTQ implementation on the Pytorch model. This gives us a ground truth of scales and
+    # Quantized values
+    with torch.no_grad():
+        weight = weights if is_3d_weights else np.expand_dims(weights, axis=0)
+        # Inputs for 3D is a list of shape [n_inputs, batch_size, 128, in_dim]
+        # For 2D it is [n_inputs, 1, 128, in_dim]. We unsqueeze at 0 for every input tensor
+        batched_inputs = inputs if is_3d_weights else [np.expand_dims(x, axis=0) for x in inputs]
+
+        pt_model = Linear3DModel(weights) if is_3d_weights else torch.nn.Linear(hidden_dims, out_dims, bias=False)
+        if not is_3d_weights:
+            pt_model.weight.copy_(torch.from_numpy(weights))
+
+        ref_gptqs, ref_scales = [], []
+        for batch in range(weight.shape[0]):
+            layer = torch.nn.Linear(hidden_dims, out_dims, bias=False)
+            layer.weight.copy_(torch.from_numpy(weight[batch]))
+
+            ref_gptq = GPTQReference(layer)
+            for inp in batched_inputs:
+                ref_gptq.add_batch(torch.from_numpy(inp[batch]))
+
+            ref_scale_for_batch, _, _ = ref_gptq.fasterquant(percdamp=0.1, group_size=group_size)
+            ref_gptqs.append(ref_gptq)
+            ref_scales.append(ref_scale_for_batch)
+
+        ref_scale = np.stack([s.detach().cpu().numpy() for s in ref_scales], axis=0)
+
+    # Step 2: Create OV models so that we can use nncf to compress this with our own GPTQ
+    # We do not use ov.convert_model() here since we expect a specific transposed weight and
+    # not transposed activation. It is hard to create and convert such a model from PT -> OV
+    # due to things like constant folding etc. which are automatically performed or generally
+    # hard to translate.
+    ov_model = _create_ov_model(weights, inputs[0].shape, is_3d_weights)
     graph = build_graph(ov_model)
 
-    # GPTQ
+    # Step 3: Setup and apply GPTQ as usual
     gptq = GPTQ()
     gptq._set_backend_entity(ov_model)
 
-    nodes = graph.get_all_nodes()
+    node_with_weight = graph.get_all_nodes()[1]
+
     wrapped_inputs = [Tensor(inp) for inp in inputs]
-    H = gptq._calculate_hessian(nodes[1], wrapped_inputs)
+    H = gptq._calculate_hessian(node_with_weight, wrapped_inputs, is_3d_weight=is_3d_weights)
 
-    ref_H = ref_gptq.H.numpy()
-    assert np.all(np.isclose(ref_H, H.data))
+    reference_gptq_list = ref_gptqs if is_3d_weights else [ref_gptq]
+    nncf_hessian_list = H.data if is_3d_weights else np.expand_dims(H.data, axis=0)
 
-    wc_params = WeightCompressionParameters(
+    for batch_index, (reference_gptq, nncf_hessian) in enumerate(zip(reference_gptq_list, nncf_hessian_list)):
+        reference_hessian = reference_gptq.H.detach().numpy()
+        assert np.all(np.isclose(reference_hessian, nncf_hessian)), f"Hessian mismatch for batch {batch_index}"
+
+    nncf_dataset = _make_nncf_dataset(ov_model, inputs)
+    reduction_axes = (1,) if not is_3d_weights else (2,)
+    wc_param = WeightCompressionParameters(
         weight_name="self.weight",
-        node_with_weight=nodes[1],
+        node_with_weight=node_with_weight,
         weight_port_id=1,
         weight_dtype=TensorDataType.float32,
         weight_shape=weights.shape,
-        reduction_axes=(1,),
+        reduction_axes=reduction_axes,
+    )
+    wc_param.compression_config = WeightCompressionConfig(mode=CompressWeightsMode.INT4_SYM, group_size=group_size)
+    wc_params = [wc_param]
+
+    _, res = gptq.apply(ov_model, graph, nncf_dataset, wc_params)
+
+    # Step 4: Obtain the scales from our GPTQ implementation and compare with reference
+    scale_from_nncf = res.get("self.weight").scale.data
+    ref_scale = ref_scale.numpy() if isinstance(ref_scale, torch.Tensor) else ref_scale
+    ref_scale = ref_scale.reshape(scale_from_nncf.shape)
+
+    assert np.all(np.isclose(ref_scale, scale_from_nncf))
+    # Here we obtain weights from the model instead of apply directly so that we also check
+    # if the weight is changed in the OV model.
+    ov_weight = gptq._backend_entity.get_weight(node_with_weight, 1, ov_model, graph)
+    ref_weights = np.stack(
+        [ref_gptq.layer.weight.detach().numpy() for ref_gptq in ref_gptqs],
+        axis=0,
     )
-    wc_params.compression_config = WeightCompressionConfig(mode=CompressWeightsMode.INT4_SYM, group_size=16)
-
-    scale, _ = gptq._quantize_weights(ov_model, graph, wc_params, H, wrapped_inputs)
-    ref_scale = ref_scale.numpy()
-    scale = scale.reshape(ref_scale.shape)
-    assert np.all(np.isclose(ref_scale, scale.data))
 
-    torch_weight = layer.weight.detach().numpy()
-    ov_weight = gptq._backend_entity.get_weight(nodes[1], 1, ov_model, graph)
-    assert np.all(np.isclose(torch_weight, ov_weight.data))
+    assert np.all(np.isclose(ref_weights, ov_weight.data))