[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

neural_compressor/adaptor/ox_utils/weight_only.py

@@ -246,6 +246,7 @@ def quant_tensor(data, num_bits=4, group_size=32, scheme="asym", dtype="int", ra
 
     return q_weight, scale, zero_point
 
+
 def quant_tensor_k_quant_cpu(data, num_bits=4, group_size=32):
     """Quantize tensor per group based on k quant.
     Ref: https://github.com/ggml-org/llama.cpp/blob/64eda5deb9859e87a020e56bab5d2f9ca956f1de/ggml/src/ggml-quants.c
@@ -260,44 +261,44 @@ def quant_tensor_k_quant_cpu(data, num_bits=4, group_size=32):
         scale: scale
         zero_point: zero point
     """
-    data = np.reshape(data, (-1, group_size)).astype(np.float32)   # (nb, group_size)
+    data = np.reshape(data, (-1, group_size)).astype(np.float32)  # (nb, group_size)
     maxq = 2**num_bits - 1
     minq = 0
-    sum_x2 = np.sum(data**2, axis=1, keepdims=True) # (nb, 1)
-    av_x = np.sqrt(sum_x2 / group_size) # (nb, 1)
-    weights = np.add(av_x, np.abs(data)) # (nb, group_size)
-    rmin = np.min(data, axis=1, keepdims=True) # (nb, 1)
-    rmax = np.max(data, axis=1, keepdims=True) # (nb, 1)
-    sum_w = np.sum(weights, axis=1, keepdims=True) # (nb, 1)
-    sum_x = np.sum(weights * data, axis=1, keepdims=True) # (nb, group_size)
-    iscale = np.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
+    sum_x2 = np.sum(data**2, axis=1, keepdims=True)  # (nb, 1)
+    av_x = np.sqrt(sum_x2 / group_size)  # (nb, 1)
+    weights = np.add(av_x, np.abs(data))  # (nb, group_size)
+    rmin = np.min(data, axis=1, keepdims=True)  # (nb, 1)
+    rmax = np.max(data, axis=1, keepdims=True)  # (nb, 1)
+    sum_w = np.sum(weights, axis=1, keepdims=True)  # (nb, 1)
+    sum_x = np.sum(weights * data, axis=1, keepdims=True)  # (nb, group_size)
+    iscale = np.ones(rmax.shape, dtype=data.dtype)  # (nb, 1)
     mask = rmin != rmax
     iscale[mask] = (maxq - minq) / (rmax[mask] - rmin[mask])
     scale = 1 / iscale
-    quant_data = np.clip(np.round(iscale * (data - rmin)), minq, maxq) # (nb, group_size)
-    diff = scale * quant_data + rmin - data # (nb, group_size)
-    best_mad = np.sum(weights * diff ** 2, axis=1, keepdims=True) # (nb, 1)
+    quant_data = np.clip(np.round(iscale * (data - rmin)), minq, maxq)  # (nb, group_size)
+    diff = scale * quant_data + rmin - data  # (nb, group_size)
+    best_mad = np.sum(weights * diff**2, axis=1, keepdims=True)  # (nb, 1)
     nstep = 20
     rdelta = 0.1
     # nstep * rdelta = -2 * rrmin, maxq - minq = 2**num_bits - 1
     rrmin = -1
     for is_ in range(nstep):
-        iscale_new = np.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
+        iscale_new = np.ones(rmax.shape, dtype=data.dtype)  # (nb, 1)
         factor = np.array([rrmin + rdelta * is_ + maxq - minq]).astype(data.dtype)[0]
         mask = rmin != rmax
         iscale_new[mask] = factor / (rmax[mask] - rmin[mask])
-        quant_data_new = np.clip(np.round(iscale_new * (data - rmin)), minq, maxq) # (nb, group_size)
+        quant_data_new = np.clip(np.round(iscale_new * (data - rmin)), minq, maxq)  # (nb, group_size)
         mul_weights_quant_data_new = weights * quant_data_new
-        sum_l = np.sum(mul_weights_quant_data_new, axis=1, keepdims=True) # (nb, 1)
-        sum_l2 = np.sum(mul_weights_quant_data_new * quant_data_new, axis=1, keepdims=True) # (nb, 1)
-        sum_xl = np.sum(mul_weights_quant_data_new * data, axis=1, keepdims=True) # (nb, 1)
-        D = np.subtract(sum_w * sum_l2, sum_l ** 2) # (nb, 1)
+        sum_l = np.sum(mul_weights_quant_data_new, axis=1, keepdims=True)  # (nb, 1)
+        sum_l2 = np.sum(mul_weights_quant_data_new * quant_data_new, axis=1, keepdims=True)  # (nb, 1)
+        sum_xl = np.sum(mul_weights_quant_data_new * data, axis=1, keepdims=True)  # (nb, 1)
+        D = np.subtract(sum_w * sum_l2, sum_l**2)  # (nb, 1)
 
-        this_scale = (sum_w * sum_xl - sum_x * sum_l) / D # (nb, 1)
-        this_min = (sum_l2 * sum_x - sum_l * sum_xl) / D # (nb, 1)
+        this_scale = (sum_w * sum_xl - sum_x * sum_l) / D  # (nb, 1)
+        this_min = (sum_l2 * sum_x - sum_l * sum_xl) / D  # (nb, 1)
 
-        diff = this_scale * quant_data_new + this_min - data # (nb, group_size)
-        mad = np.sum(weights * diff ** 2, axis=1, keepdims=True) # (nb, 1)
+        diff = this_scale * quant_data_new + this_min - data  # (nb, group_size)
+        mad = np.sum(weights * diff**2, axis=1, keepdims=True)  # (nb, 1)
 
         mad_1 = np.array(mad)
         best_mad_1 = np.array(best_mad)
@@ -307,7 +308,7 @@ def quant_tensor_k_quant_cpu(data, num_bits=4, group_size=32):
         scale[idx_to_replace] = this_scale[idx_to_replace]
         rmin[idx_to_replace] = this_min[idx_to_replace]
 
-    zero_point = np.clip((( - rmin) / scale).round(), 0, maxq).astype("uint8")
+    zero_point = np.clip(((-rmin) / scale).round(), 0, maxq).astype("uint8")
     scale = scale.astype(np.float64)
     q_weight = np.empty_like(data, dtype=scale.dtype)
     np.divide(data, scale, out=q_weight)
@@ -317,6 +318,7 @@ def quant_tensor_k_quant_cpu(data, num_bits=4, group_size=32):
 
     return q_weight, scale, zero_point
 
+
 def quant_tensor_k_quant_cuda(data, num_bits=4, group_size=32):
     """Quantize tensor per group based on k quant.
     Ref: https://github.com/ggml-org/llama.cpp/blob/64eda5deb9859e87a020e56bab5d2f9ca956f1de/ggml/src/ggml-quants.c
@@ -334,46 +336,47 @@ def quant_tensor_k_quant_cuda(data, num_bits=4, group_size=32):
     try:
         import cupy as cp
         import torch
+
         if torch.cuda.is_available():
             data = cp.asarray(data)
-            data = data.reshape((-1, group_size)).astype(np.float32)   # (nb, group_size)
+            data = data.reshape((-1, group_size)).astype(np.float32)  # (nb, group_size)
             nb = data.shape[0]
             maxq = 2**num_bits - 1
             minq = 0
-            sum_x2 = np.sum(data**2, axis=1, keepdims=True) # (nb, 1)
-            av_x = np.sqrt(sum_x2 / group_size) # (nb, 1)
-            weights = np.add(av_x, np.abs(data)) # (nb, group_size)
-            rmin = np.min(data, axis=1, keepdims=True) # (nb, 1)
-            rmax = np.max(data, axis=1, keepdims=True) # (nb, 1)
-            sum_w = np.sum(weights, axis=1, keepdims=True) # (nb, 1)
-            sum_x = np.sum(weights * data, axis=1, keepdims=True) # (nb, group_size)
-            iscale = cp.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
+            sum_x2 = np.sum(data**2, axis=1, keepdims=True)  # (nb, 1)
+            av_x = np.sqrt(sum_x2 / group_size)  # (nb, 1)
+            weights = np.add(av_x, np.abs(data))  # (nb, group_size)
+            rmin = np.min(data, axis=1, keepdims=True)  # (nb, 1)
+            rmax = np.max(data, axis=1, keepdims=True)  # (nb, 1)
+            sum_w = np.sum(weights, axis=1, keepdims=True)  # (nb, 1)
+            sum_x = np.sum(weights * data, axis=1, keepdims=True)  # (nb, group_size)
+            iscale = cp.ones(rmax.shape, dtype=data.dtype)  # (nb, 1)
             mask = rmin != rmax
             iscale[mask] = (maxq - minq) / (rmax[mask] - rmin[mask])
             scale = 1 / iscale
-            quant_data = np.clip(np.round(iscale * (data - rmin)), minq, maxq) # (nb, group_size)
-            diff = scale * quant_data + rmin - data # (nb, group_size)
-            best_mad = np.sum(weights * diff ** 2, axis=1, keepdims=True) # (nb, 1)
+            quant_data = np.clip(np.round(iscale * (data - rmin)), minq, maxq)  # (nb, group_size)
+            diff = scale * quant_data + rmin - data  # (nb, group_size)
+            best_mad = np.sum(weights * diff**2, axis=1, keepdims=True)  # (nb, 1)
             nstep = 20
             rdelta = 0.1
             rrmin = -1
             for is_ in range(nstep):
-                iscale_new = cp.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
+                iscale_new = cp.ones(rmax.shape, dtype=data.dtype)  # (nb, 1)
                 factor = cp.array([rrmin + rdelta * is_ + maxq - minq]).astype(data.dtype)[0]
                 mask = rmin != rmax
                 iscale_new[mask] = factor / (rmax[mask] - rmin[mask])
-                quant_data_new = np.clip(np.round(iscale_new * (data - rmin)), minq, maxq) # (nb, group_size)
+                quant_data_new = np.clip(np.round(iscale_new * (data - rmin)), minq, maxq)  # (nb, group_size)
                 mul_weights_quant_data_new = weights * quant_data_new
-                sum_l = np.sum(mul_weights_quant_data_new, axis=1, keepdims=True) # (nb, 1)
-                sum_l2 = np.sum(mul_weights_quant_data_new * quant_data_new, axis=1, keepdims=True) # (nb, 1)
-                sum_xl = np.sum(mul_weights_quant_data_new * data, axis=1, keepdims=True) # (nb, 1)
-                D = np.subtract(sum_w * sum_l2, sum_l ** 2) # (nb, 1)
+                sum_l = np.sum(mul_weights_quant_data_new, axis=1, keepdims=True)  # (nb, 1)
+                sum_l2 = np.sum(mul_weights_quant_data_new * quant_data_new, axis=1, keepdims=True)  # (nb, 1)
+                sum_xl = np.sum(mul_weights_quant_data_new * data, axis=1, keepdims=True)  # (nb, 1)
+                D = np.subtract(sum_w * sum_l2, sum_l**2)  # (nb, 1)
 
-                this_scale = (sum_w * sum_xl - sum_x * sum_l) / D # (nb, 1)
-                this_min = (sum_l2 * sum_x - sum_l * sum_xl) / D # (nb, 1)
+                this_scale = (sum_w * sum_xl - sum_x * sum_l) / D  # (nb, 1)
+                this_min = (sum_l2 * sum_x - sum_l * sum_xl) / D  # (nb, 1)
 
-                diff = this_scale * quant_data_new + this_min - data # (nb, group_size)
-                mad = np.sum(weights * diff ** 2, axis=1, keepdims=True) # (nb, 1)
+                diff = this_scale * quant_data_new + this_min - data  # (nb, group_size)
+                mad = np.sum(weights * diff**2, axis=1, keepdims=True)  # (nb, 1)
 
                 mad_1 = cp.array(mad)
                 best_mad_1 = cp.array(best_mad)
@@ -383,7 +386,7 @@ def quant_tensor_k_quant_cuda(data, num_bits=4, group_size=32):
                 scale[idx_to_replace] = this_scale[idx_to_replace]
                 rmin[idx_to_replace] = this_min[idx_to_replace]
 
-            zero_point = np.clip((( - rmin) / scale).round(), 0, maxq).astype("uint8")
+            zero_point = np.clip(((-rmin) / scale).round(), 0, maxq).astype("uint8")
             scale = scale.astype(np.float64)
             q_weight = np.empty_like(data, dtype=scale.dtype)
             np.divide(data, scale, out=q_weight)
@@ -393,20 +396,18 @@ def quant_tensor_k_quant_cuda(data, num_bits=4, group_size=32):
 
             return q_weight.get(), scale.get(), zero_point.get()
         else:
-            logger.warning("Try to use k-quant quantization on CUDA. However, CUDA is not available." \
-                           "Fall back to k-quant quantization on CPU.")
-            return quant_tensor_k_quant_cpu(
-                data, num_bits, group_size
+            logger.warning(
+                "Try to use k-quant quantization on CUDA. However, CUDA is not available."
+                "Fall back to k-quant quantization on CPU."
             )
+            return quant_tensor_k_quant_cpu(data, num_bits, group_size)
     except ImportError:
         logger.info(
-            "Now we are using k-quant quantization on cpu, which is time consuming." \
-            "Please consider install cupy to speed up on CUDA. See https://cupy.dev/" \
-            "Please also install torch to check CUDA availablity."
-            )
-        return quant_tensor_k_quant_cpu(
-            data, num_bits, group_size
+            "Now we are using k-quant quantization on cpu, which is time consuming."
+            "Please consider install cupy to speed up on CUDA. See https://cupy.dev/"
+            "Please also install torch to check CUDA availability."
         )
+        return quant_tensor_k_quant_cpu(data, num_bits, group_size)
 
 
 def qdq_tensor(data, num_bits=4, group_size=32, scheme="asym", dtype="int", ratio=1.0):
@@ -536,9 +537,7 @@ def rtn_quantize(
                 # MatMulFpQ4 support 4 bits and 32 group_size with ort 1.16.0 and 1.16.1 versions, supported by CPU EP
                 # MatMulNBits supports 4 bits and 2^n group_size with ort > 1.16.1, supported by CPU EP AND CUDA EP
                 if algorithm == "k_quant":
-                    q_weight, scale, zp = quant_tensor_k_quant_cuda(
-                        weight.T, num_bits, group_size
-                    )
+                    q_weight, scale, zp = quant_tensor_k_quant_cuda(weight.T, num_bits, group_size)
                 else:
                     q_weight, scale, zp = quant_tensor(
                         weight.T, num_bits, group_size, scheme, "uint", ratios.get(node.input[1], 1)