fix prev weights case

Author: anzr299

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

@@ -218,13 +218,7 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
         X = X.astype(TensorDataType.float32)
 
         is_3d_weight = weight.ndim == 3
-        if s.ndim == 1:
-            s = fns.unsqueeze(s, 0)  # [hidden_dim] -> [1, hidden_dim]
-            X = fns.unsqueeze(X, 0)  # [hidden_dim, samples] -> [1, hidden_dim, samples]
-            weight = fns.unsqueeze(weight, 0)  # [out_features, hidden_dim] -> [1, out_features, hidden_dim]
-            is_2d_weight = True
-        else:
-            is_2d_weight = False
+        is_2d_weight = weight.ndim == 2
 
         assert isinstance(wp.reduction_axes, tuple) and len(wp.reduction_axes) == 1
         reduction_axis = wp.reduction_axes[0]
@@ -233,7 +227,13 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
         if prev_statistics is not None and prev_weight is not None:
             prev_s, _ = process_stats(prev_statistics, self._subset_size)
             prev_s = prev_s.astype(TensorDataType.float32).max().item()
-            prev_w = fns.mean(fns.abs(prev_weight), axis=reduction_axis)
+            prev_weight = fns.unsqueeze(prev_weight, 0)  # [out_features, hidden_dim] -> [1, out_features, hidden_dim]
+            prev_w = fns.mean(fns.abs(prev_weight), axis=reduction_axis + 1)
+
+        if is_2d_weight:
+            s = fns.unsqueeze(s, 0)  # [hidden_dim] -> [1, hidden_dim]
+            X = fns.unsqueeze(X, 0)  # [hidden_dim, samples] -> [1, hidden_dim, samples]
+            weight = fns.unsqueeze(weight, 0)  # [out_features, hidden_dim] -> [1, out_features, hidden_dim]
 
         top_k = max(int(s.shape[-1] * self._percent_to_apply), 1)
         topk_idxs = fns.argsort(-s)[:, :top_k]
@@ -289,7 +289,9 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
                     # per channel magnitudes for the previous MatMul
                     # mean(abs(prev_weight)) * max(abs((prev_activation))) * prev_weight.shape[reduction_axis]
                     magnitudes = (
-                        (prev_w[offset : offset + group_size] / cur_scale) * prev_s * prev_weight.shape[reduction_axis]
+                        (prev_w[expert_idx, offset : offset + group_size] / cur_scale)
+                        * prev_s
+                        * prev_weight.shape[reduction_axis]
                     )
                     if magnitudes.max() >= threshold:
                         cur_scale = AWQ._clamp_scale(