Repository: quic/aimet (language: python)

Issue:
Add blockwise quant with batched matmul implementation



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_batched_matmul.py b/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_batched_matmul.py
new file mode 100644
index 0000000000..b597b07a2a
--- /dev/null
+++ b/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_batched_matmul.py
@@ -0,0 +1,367 @@
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+""" Utilities for implementing blockwise quantization using batched matmul approach """
+
+import functools
+from typing import List, Union
+import torch
+
+import aimet_common.AimetTensorQuantizer as AimetTensorQuantizer
+import aimet_common.libpymo as libpymo
+from aimet_common.defs import QuantScheme, QuantizationDataType, MAP_QUANT_SCHEME_TO_PYMO
+from aimet_torch import utils
+from aimet_torch.qc_quantize_op import QcQuantizeWrapper
+from aimet_torch.quantsim import QuantizationSimModel
+from aimet_torch.quantsim_straight_through_grad import get_computed_encodings
+from aimet_torch.tensor_quantizer import StaticGridTensorQuantizer, StaticGridPerChannelQuantizer
+
+
+class ReduceSum(torch.nn.Module):
+    """
+    ReduceSum implementation
+    """
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+    def forward(self, x):
+        """
+        Forward pass.
+
+        :param x: Input tensor
+        :return: Output of forward pas
+        """
+        return torch.sum(x, dim=self.dim)
+
+class BatchMatMulWithWeight(torch.nn.Module):
+    """
+    BatchMatMulWithWeight implementation
+    """
+    def __init__(self, weight: torch.Tensor, weight_first: bool):
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight)
+        self.weight_first = weight_first
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        :param x: Input tensor
+        :return: Output of forward pas
+        """
+        if self.weight_first:
+            return torch.matmul(self.weight, x)
+        return torch.matmul(x, self.weight)
+
+
+class LinearWithMatMul(torch.nn.Module):
+    '''
+    Linear(out_feature, in_feature)
+    x:[n, seq, in_feature] * w.T:[in_feature, out_feature] = [n, seq, out_feature]
+    where, in_feature = num_blocks * block_size
+    Approach #1: weight is the 1st input of torch.matmul
+    w:[out_feature, in_feature].reshape()
+    => [out_feature, num_blocks, block_size].permute(1, 0, 2)
+    => [num_blocks, out_feature, block_size]
+    x:[n, seq, in_feature].reshape()
+    => [n * seq, num_blocks, block_size].permute(1, 2, 0)
+    => [num_blocks, block_size, n * seq]
+    torch.matmul(w, x)
+    => [num_blocks, out_feature, n * seq]
+    => [num_blocks, out_feature, n * seq].sum(dim=0)
+    => [out_feature, n * seq].reshape()
+    => [out_feature, n, seq].permute(1, 2, 0)
+    Approach #2: weight is the 2nd input of torch.matmul
+    w:[out_feature, in_feature].reshape()
+    => [out_feature, num_blocks, block_size].permute(1, 2, 0)
+    => [num_blocks, block_size, out_feature]
+    x:[n, seq, in_feature].reshape()
+    => [n * seq, num_blocks, block_size].permute(1, 0, 2)
+    => [num_blocks, n * seq, block_size]
+    torch.matmul(x, w)
+    => [num_blocks, n * seq, out_feature].sum(dim=0)
+    => [1, n* seq, out_feature].reshape()
+    => [n, seq, out_feature]
+    '''
+    def __init__(self, linear: torch.nn.Module, block_size: int, weight_first=True):
+        super().__init__()
+        _, in_feature = linear.weight.shape
+        assert in_feature % block_size == 0, f"in_feature:{in_feature} should be divisible by block_size:{block_size}"
+
+        self.block_size = block_size
+        self.num_blocks = in_feature // block_size
+
+        self.weight_first = weight_first
+        if self.weight_first:
+            weight = linear.weight.detach().reshape(-1, self.num_blocks, block_size).permute(1, 0, 2)
+        else:
+            weight = linear.weight.detach().reshape(-1, self.num_blocks, block_size).permute(1, 2, 0)
+
+        self.batch_matmul = BatchMatMulWithWeight(weight, weight_first=self.weight_first)
+        self.reducesum = ReduceSum(dim=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        :param x: Input tensor
+        :return: Output of forward pas
+        """
+
+        batch, seq = x.shape[:2]
+
+        if self.weight_first:
+            x = x.reshape(batch * seq, self.num_blocks, self.block_size).permute(1, 2, 0)
+            x = self.batch_matmul(x)
+            x = self.reducesum(x).reshape(batch, seq, -1)
+        else:
+            x = x.reshape(batch * seq, self.num_blocks, self.block_size).permute(1, 0, 2)
+            x = self.batch_matmul(x)
+            x = self.reducesum(x).reshape(batch, seq, -1)
+        return x
+
+class StaticGridMultiDimPerChannelQuantizer(StaticGridTensorQuantizer):
+    """
+    Quantizer allowing multiple channels to be specified as channel axes. The number of channels will be the product of
+    sizes of all channels specified as channel axes.
+    """
+    # pylint: disable=too-many-arguments
+    def __init__(self, bitwidth: int, round_mode: libpymo.RoundingMode, quant_scheme: QuantScheme,
+                 use_symmetric_encodings: bool, num_channels: int, enabled_by_default: bool, ch_axes: Union[int, list],
+                 data_type: QuantizationDataType = QuantizationDataType.int):
+        super(StaticGridMultiDimPerChannelQuantizer, self).__init__(
+            bitwidth, round_mode, quant_scheme, use_symmetric_encodings, enabled_by_default, data_type=data_type)
+
+        quant_scheme = MAP_QUANT_SCHEME_TO_PYMO[quant_scheme]
+        self._cppOp = [AimetTensorQuantizer.AimetTensorQuantizer(quant_scheme) for _ in range(num_channels)]
+        if not isinstance(ch_axes, list):
+            self._ch_axes = [ch_axes]
+        else:
+            self._ch_axes = ch_axes
+        self._permute_order = None
+        self._encoding_reshape_order = None
+
+    @property
+    def encoding(self) -> Union[None, List[libpymo.TfEncoding]]:
+        """
+        Property to get encoding.
+
+        :return: List of Encoding(s).
+        """
+        return self._encoding
+
+    @encoding.setter
+    def encoding(self, encoding: List[libpymo.TfEncoding]):
+        """
+        Property to set encoding.
+
+        :param encoding: List of Encoding(s).
+        """
+        if self._is_encoding_frozen:
+            raise RuntimeError("Encoding can be set only when it is not frozen.")
+
+        self._encoding = encoding
+
+    def _get_permute_order(self, num_dims: int) -> List[int]:
+        """
+        Return the permute order for permuting all channel axes to the front dimensions.
+
+        :param num_dims: Number of dimensions of tensor to permute
+        :return: Tensor permute order
+        """
+        if self._permute_order is not None:
+            return self._permute_order
+
+        remaining_dims = set(range(num_dims))
+        permute_order = []
+        for ch_axis in self._ch_axes:
+            assert ch_axis < num_dims
+            permute_order.append(ch_axis)
+            remaining_dims.remove(ch_axis)
+        for i in range(num_dims):
+            if i in remaining_dims:
+                permute_order.append(i)
+        self._permute_order = permute_order
+        return self._permute_order
+
+    def _get_encoding_reshape_order(self, tensor_shape: torch.Size) -> List[int]:
+        """
+        Return the reshape order encodings should be reshaped to. Indices corresponding to channel axes will have proper
+        sizes and all remaining indices will be of size 1.
+
+        :param tensor_shape: Shape of tensor to quantize
+        :return: Encoding reshape order
+        """
+        if self._encoding_reshape_order is not None:
+            return self._encoding_reshape_order
+
+        encoding_reshape_order = []
+        for idx, _ in enumerate(tensor_shape):
+            if idx in self._ch_axes:
+                encoding_reshape_order.append(tensor_shape[idx])
+            else:
+                encoding_reshape_order.append(1)
+        self._encoding_reshape_order = encoding_reshape_order
+        return self._encoding_reshape_order
+
+    def update_encoding_stats(self, tensor: torch.Tensor):
+        """
+        Update the stats for computing encodings. In the case of fp16 or int with bw=32, this is skipped.
+
+        :param tensor: Tensor to use for updating the encodings stats
+        """
+        if self.enabled and not self._is_encoding_frozen:
+            if self.bitwidth == 32:
+                return
+            if self.data_type == QuantizationDataType.float:
+                raise AssertionError('Float quantization not supported')
+
+            if self.encoding_min_max_fixed_vals is not None:
+                # pylint: disable=unsubscriptable-object
+                tensor = torch.tensor([self.encoding_min_max_fixed_vals[0],
+                                       self.encoding_min_max_fixed_vals[1]])
+
+                for op in self._cppOp:
+                    op.updateStats(tensor, tensor.is_cuda)
+            else:
+                # If we have a half-float tensor, just upcast it to float
+                # Need to change this when libpymo can handle half-float tensors
+                if tensor.dtype == torch.float16:
+                    tensor = tensor.to(torch.float32)
+
+                # Permute tensor so all channels to coalesce are at the front dimensions
+                tensor_slice = torch.permute(tensor, dims=self._get_permute_order(len(tensor.shape)))
+
+                # Reshape tensor to combine all channels to coalesce into a single dimension; all remaining axes
+                # are also coalesced into a second dimension
+                tensor_slice = torch.reshape(tensor_slice,
+                                             (functools.reduce(lambda x, y: x * y,
+                                                               tensor_slice.shape[:len(self._ch_axes)]),
+                                              -1))
+                for channel_idx, op in enumerate(self._cppOp):
+                    channel_slice = tensor_slice.select(0, channel_idx).contiguous(
+                        memory_format=torch.contiguous_format)
+                    op.updateStats(channel_slice, tensor.is_cuda)
+
+    def quantize_dequantize(self, tensor: torch.Tensor, round_mode: libpymo.RoundingMode) -> torch.Tensor:
+        """
+        Quantize-dequantize the tensor, using the saved encoding for this tensor
+
+        :param tensor: Tensor to quantize-dequantize
+        :param round_mode: Rounding mode
+        :return: Resulting tensor
+        """
+        encoding_mins = [encoding.min for encoding in self._encoding]
+        encoding_maxes = [encoding.max for encoding in self._encoding]
+        encoding_min_tensor = torch.tensor(encoding_mins, device=tensor.device)
+        encoding_max_tensor = torch.tensor(encoding_maxes, device=tensor.device)
+        delta, offset, num_steps = get_computed_encodings(self.bitwidth, encoding_min_tensor, encoding_max_tensor,
+                                                          self.use_symmetric_encodings, self.use_strict_symmetric,
+                                                          self.is_unsigned_symmetric)
+        offset = offset.expand_as(delta)     # In case of symmetric encodings, offset will be a scalar before expansion
+        delta = torch.reshape(delta, self._get_encoding_reshape_order(tensor.shape))
+        offset = torch.reshape(offset, self._get_encoding_reshape_order(tensor.shape))
+
+        zero = torch.zeros_like(num_steps)
+
+        x_round = torch.round(tensor / delta) - offset
+        x_quant = x_round.clamp(zero, num_steps)
+        return (x_quant + offset) * delta
+
+def replace_linears_for_blockwise_quant(model: torch.nn.Module, block_size: int):
+    """
+    Replace all instances of torch.nn.Linears in model with equivalent LinearWithMatMul modules.
+
+    :param model: Model to replace nn.Linears for
+    :param block_size: Block size to use
+    """
+    linear_layers = []
+    for name, module in model.named_children():
+        if isinstance(module, torch.nn.Linear):
+            linear_layers.append((name, module))
+        elif not utils.is_leaf_module(module):
+            replace_linears_for_blockwise_quant(module, block_size)
+
+    for name, module in linear_layers:
+        setattr(model, name, LinearWithMatMul(module, block_size))
+
+def enable_bmm_per_channel_quantizers(quantsim: QuantizationSimModel):
+    """
+    Loop through quantsim modules and replace weight param quantizers of BatchMatMulWithWeight modules with multi
+    channel quantizers.
+
+    :param quantsim: Quantsim containing BatchMatMulWithWeight modules
+    """
+    for module in quantsim.model.modules():
+        # pylint: disable=protected-access
+        if isinstance(module, QcQuantizeWrapper) and isinstance(module._module_to_wrap, BatchMatMulWithWeight):
+            if isinstance(module.param_quantizers['weight'], StaticGridMultiDimPerChannelQuantizer):
+                continue
+
+            if not isinstance(module.param_quantizers['weight'], StaticGridPerChannelQuantizer):
+                module.enable_per_channel_quantization()
+
+            # pylint: disable=protected-access
+            module.param_quantizers['weight'] = create_multi_channel_quantizer(module.param_quantizers['weight'],
+                                                                               module._module_to_wrap.weight_first,
+                                                                               module._module_to_wrap.weight.shape)
+
+def create_multi_channel_quantizer(per_channel_quantizer: StaticGridPerChannelQuantizer, weight_first: bool,
+                                   weight_shape: torch.Size) -> StaticGridMultiDimPerChannelQuantizer:
+    """
+    Create multi channel quantizer given an original per channel quantizer.
+
+    :param per_channel_quantizer: Original per channel quantizer to obtain quantizer settings from
+    :param weight_first: True if the LinearWithMatMul is using weight first, False otherwise
+    :param weight_shape: Shape of weight parameter
+    :return: Multi channel quantizer
+    """
+    if weight_first:
+        channel_axes = [0, 1]
+        num_channels = weight_shape[0] * weight_shape[1]
+    else:
+        channel_axes = [0, 2]
+        num_channels = weight_shape[0] * weight_shape[2]
+
+    return StaticGridMultiDimPerChannelQuantizer(per_channel_quantizer.bitwidth,
+                                                 per_channel_quantizer.round_mode,
+                                                 per_channel_quantizer.quant_scheme,
+                                                 per_channel_quantizer.use_symmetric_encodings,
+                                                 num_channels,
+                                                 per_channel_quantizer.enabled,
+                                                 channel_axes,
+                                                 per_channel_quantizer.data_type
+                                                 )
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_grouped_conv.py b/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_grouped_conv.py
new file mode 100644
index 0000000000..d5aa7b65a4
--- /dev/null
+++ b/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_grouped_conv.py
@@ -0,0 +1,129 @@
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+""" Utilities for implementing blockwise quantization using grouped conv approach """
+
+import torch
+from aimet_torch import utils
+
+class ReduceSum(torch.nn.Module):
+    """
+    ReduceSum implementation
+    """
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+    def forward(self, x):
+        """
+        Forward pass.
+
+        :param x: Input tensor
+        :return: Output of forward pas
+        """
+        return torch.sum(x, dim=self.dim)
+
+class LinearWithGroupedConv(torch.nn.Module):
+    '''
+    Linear(out_feature, in_feature)
+    x:[n, seq, in_feature] * w.T:[in_feature, out_feature] = [n, seq, out_feature]
+    where, in_feature = num_blocks * block_size
+    w:[out_feature, in_feature].reshape()
+    => [out_feature, num_blocks, block_size].permute(1,0,2)
+    => [num_blocks, out_feature, block_size].reshape()
+    => [num_blocks * out_feature, block_size, 1, 1]
+    x:[n, seq, in_feature].permute(0, 2, 1)
+    => [n, in_feature, seq].reshape()
+    => [n, in_feature, 1, seq]
+    => torch.conv(weight=w, input=x, groups=num_blocks)
+    => [n, num_blocks * out_feature, 1, seq].reshape()
+    => [n, num_blocks, out_feature, seq].sum(dim=1)
+    => [n, out_feature, seq].permute(0,2,1)
+    [n, seq, out_feature]
+    '''
+    def __init__(self, linear: torch.nn.Module, block_size: int):
+        super().__init__()
+        out_feature, in_feature = linear.weight.shape
+        assert in_feature % block_size == 0, f"in_feature:{in_feature} should be divisible by block_size:{block_size}"
+
+        self.block_size = block_size
+        self.num_blocks = in_feature // block_size
+
+        weight = linear.weight.detach().reshape(-1, self.num_blocks, self.block_size)
+        weight = weight.permute(1, 0, 2)
+        weight = weight.reshape(-1, self.block_size, 1, 1)
+
+        self.grouped_conv = torch.nn.Conv2d(in_channels=in_feature,
+                                            out_channels=out_feature * self.num_blocks,
+                                            kernel_size=1,
+                                            groups=self.num_blocks,
+                                            bias=False)
+        self.grouped_conv.weight.data.copy_(weight)
+        self.reducesum = ReduceSum(dim=1)
+
+
+    def forward(self, x: torch.Tensor):
+        """
+        Forward pass.
+
+        :param x: Input tensor
+        :return: Output of forward pas
+        """
+        batch, seq, in_feature = x.shape
+        x = x.permute(0, 2, 1)
+        x = x.reshape(batch, in_feature, 1, seq)
+        x = self.grouped_conv(x)
+        x = x.reshape(batch, self.num_blocks, -1, seq)
+        x = self.reducesum(x)
+        x = x.permute(0, 2, 1)
+
+        return x
+
+def replace_linears_for_blockwise_quant(model: torch.nn.Module, block_size: int):
+    """
+    Replace all instances of torch.nn.Linears in model with equivalent LinearWithMatMul modules.
+
+    :param model: Model to replace nn.Linears for
+    :param block_size: Block size to use
+    """
+    linear_layers = []
+    for name, module in model.named_children():
+        if isinstance(module, torch.nn.Linear):
+            linear_layers.append((name, module))
+        elif not utils.is_leaf_module(module):
+            replace_linears_for_blockwise_quant(module, block_size)
+
+    for name, module in linear_layers:
+        setattr(model, name, LinearWithGroupedConv(module, block_size))
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_tensor_split.py b/TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_tensor_split.py
similarity index 100%
rename from TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_tensor_split.py
rename to TrainingExtensions/torch/src/python/aimet_torch/blockwise_quant_recipes/blockwise_quant_tensor_split.py


Test patch:
diff --git a/TrainingExtensions/torch/test/python/test_blockwise_quant_recipes.py b/TrainingExtensions/torch/test/python/test_blockwise_quant_recipes.py
new file mode 100644
index 0000000000..8a7bcd497d
--- /dev/null
+++ b/TrainingExtensions/torch/test/python/test_blockwise_quant_recipes.py
@@ -0,0 +1,309 @@
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+""" Tests for blockwise quant recipes """
+import copy
+import json
+import itertools
+import pytest
+import torch
+from aimet_torch.quantsim import QuantizationSimModel
+from aimet_torch.qc_quantize_op import QcQuantizeWrapper
+from aimet_torch.blockwise_quant_recipes import blockwise_quant_batched_matmul as bqbm
+from aimet_torch.blockwise_quant_recipes import blockwise_quant_grouped_conv as bqgc
+from aimet_torch.blockwise_quant_recipes import blockwise_quant_tensor_split as bqts
+
+class LinearModel(torch.nn.Module):
+    def __init__(self):
+        super(LinearModel, self).__init__()
+        self.linear1 = torch.nn.Linear(9, 3)
+        self.relu1 = torch.nn.ReLU()
+        self.linear2 = torch.nn.Linear(3, 2, bias=False)
+        self.softmax = torch.nn.Softmax()
+
+    def forward(self, inp):
+        x = self.linear1(inp)
+        x = self.relu1(x)
+        x = self.linear2(x)
+        x = self.softmax(x)
+        return x
+
+class InnerModel(torch.nn.Module):
+    def __init__(self):
+        super(InnerModel, self).__init__()
+        self.linear1 = torch.nn.Linear(8, 3)
+        self.seq = torch.nn.Sequential(torch.nn.Linear(3, 4), torch.nn.Linear(4, 5), torch.nn.ReLU())
+        self.seq2 = torch.nn.Sequential(torch.nn.Linear(5, 2))
+
+    def forward(self, inp):
+        x = self.linear1(inp)
+        x = self.seq(x)
+        x = self.seq2(x)
+        return x
+
+class NestedModel(torch.nn.Module):
+    def __init__(self):
+        super(NestedModel, self).__init__()
+        self.sigmoid = torch.nn.Sigmoid()
+        self.inner_linears = InnerModel()
+
+    def forward(self, inp):
+        x = self.inner_linears(inp)
+        x = self.sigmoid(x)
+        return x
+
+@pytest.fixture(scope='module')
+def create_per_channel_config():
+    quantsim_config = {
+        "defaults": {
+            "ops": {
+                "is_output_quantized": "True",
+                "is_symmetric": "False"
+            },
+            "params": {
+                "is_quantized": "True",
+                "is_symmetric": "True"
+            },
+            "per_channel_quantization": "True",
+        },
+        "params": {},
+        "op_type": {
+            'Split': {
+                'is_output_quantized': "False"
+            }
+        },
+        "supergroups": [],
+        "model_input": {},
+        "model_output": {}
+    }
+    with open('./data/quantsim_config.json', 'w') as f:
+        json.dump(quantsim_config, f)
+
+class TestBlockwiseQuantTensorSplitting:
+    @pytest.mark.parametrize('model, dummy_input, block_size', [(torch.nn.Linear(8, 3), torch.randn(1, 8), 3),
+                                                                 (torch.nn.Linear(8, 3, bias=False), torch.randn(1, 8), 3),
+                                                                 (torch.nn.Linear(3, 2), torch.randn(1, 3), 3),
+                                                                 (torch.nn.Linear(3, 2), torch.randn(1, 3), 4)])
+    def test_blockwise_linear(self, model, dummy_input, block_size):
+        blockwise_linear = bqts.BlockwiseLinear(model, block_size=block_size)
+        orig_out = model(dummy_input)
+        new_out = blockwise_linear(dummy_input)
+        assert torch.allclose(orig_out, new_out, atol=1e-6)
+
+    def test_replace_linears_for_blockwise_quant(self):
+        dummy_input = torch.randn(1, 9)
+        model = LinearModel()
+        linear1 = model.linear1
+        orig_out = model(dummy_input)
+        bqts.replace_linears_for_blockwise_quant(model, 3)
+
+        assert len(model.linear1.linears) == 3
+        assert torch.equal(model.linear1.linears[0].bias, linear1.bias)
+        for linear in model.linear1.linears[1:]:
+            assert linear.bias is None
+
+        assert len(model.linear1.elementwise_adds) == 2
+        assert len(model.linear2.linears) == 1
+        assert model.linear2.elementwise_adds is None
+        new_out = model(dummy_input)
+        assert torch.allclose(orig_out, new_out, atol=1e-6)
+
+    def test_quantize_blockwise_linear(self, create_per_channel_config):
+        dummy_input = torch.randn(1, 8)
+        model = bqts.BlockwiseLinear(torch.nn.Linear(8, 3), 3)
+        qsim = QuantizationSimModel(model, dummy_input=dummy_input, config_file='./data/quantsim_config.json')
+
+        assert isinstance(qsim.model.split, QcQuantizeWrapper)
+        assert isinstance(qsim.model.linears[0], QcQuantizeWrapper)
+        assert isinstance(qsim.model.linears[1], QcQuantizeWrapper)
+        assert isinstance(qsim.model.linears[2], QcQuantizeWrapper)
+        assert isinstance(qsim.model.elementwise_adds[0], QcQuantizeWrapper)
+        assert isinstance(qsim.model.elementwise_adds[1], QcQuantizeWrapper)
+
+        # Temporary hack to disable split op output quantizers while handling for CG split op is reworked
+        for output_quantizer in qsim.model.split.output_quantizers:
+            output_quantizer.enabled = False
+
+        # Temporary hack to enable model input split op input quantizer while handling for CG split op is reworked
+        qsim.model.split.input_quantizers[0].enabled = True
+
+        qsim.compute_encodings(lambda m, _: m(dummy_input), None)
+        bqts.tie_blockwise_linear_quantizers(qsim)
+        _ = qsim.model(dummy_input)
+
+        assert len(qsim.model.linears[0].param_quantizers['weight'].encoding) == 3
+        assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
+                qsim.model.linears[1].output_quantizers[0].encoding.max)
+        assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
+                qsim.model.linears[2].output_quantizers[0].encoding.max)
+        assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
+                qsim.model.elementwise_adds[0].output_quantizers[0].encoding.max)
+        assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
+                qsim.model.elementwise_adds[1].output_quantizers[0].encoding.max)
+
+    def test_blockwise_quant_with_small_linear(self, create_per_channel_config):
+        dummy_input = torch.randn(1, 3)
+        model = bqts.BlockwiseLinear(torch.nn.Linear(3, 2), 3)
+        qsim = QuantizationSimModel(model, dummy_input=dummy_input)
+        # Temporary hack to disable split op output quantizers while handling for CG split op is reworked
+        for output_quantizer in qsim.model.split.output_quantizers:
+            output_quantizer.enabled = False
+
+        # Temporary hack to enable model input split op input quantizer while handling for CG split op is reworked
+        qsim.model.split.input_quantizers[0].enabled = True
+
+        bqts.tie_blockwise_linear_quantizers(qsim)
+        qsim.compute_encodings(lambda m, _: m(dummy_input), None)
+        assert len(qsim.model.linears) == 1
+        assert qsim.model.elementwise_adds is None
+
+        _ = qsim.model(dummy_input)
+        assert len(qsim.connected_graph.get_all_ops()) == 1
+
+    def test_nested_sequential_linears(self):
+        model = NestedModel()
+        bqts.replace_linears_for_blockwise_quant(model, block_size=2)
+        num_blockwise_linears = 0
+        for _, module in model.named_modules():
+            if isinstance(module, bqts.BlockwiseLinear):
+                num_blockwise_linears += 1
+        assert num_blockwise_linears == 4
+
+class TestBlockwiseQuantBatchedMatmul:
+    def test_blockwise_quant_with_batched_matmul(self, create_per_channel_config):
+        torch.manual_seed(0)
+        dummy_input = torch.randn(1, 1, 9)
+        model = LinearModel()
+        orig_out = model(dummy_input)
+        orig_linear = model.linear1
+
+        model.linear1 = bqbm.LinearWithMatMul(orig_linear, 3)
+        new_out = model(dummy_input)
+
+        assert torch.allclose(orig_out, new_out, atol=1e-5)
+
+        qsim1 = QuantizationSimModel(model, dummy_input, config_file='./data/quantsim_config.json')
+        bqbm.enable_bmm_per_channel_quantizers(qsim1)
+        qsim1.compute_encodings(lambda m, _: m(dummy_input), None)
+        weight_first_quantized_out = qsim1.model(dummy_input)
+
+        model.linear1 = bqbm.LinearWithMatMul(orig_linear, 3, weight_first=False)
+        qsim2 = QuantizationSimModel(model, dummy_input, config_file='./data/quantsim_config.json')
+        bqbm.enable_bmm_per_channel_quantizers(qsim2)
+        qsim2.compute_encodings(lambda m, _: m(dummy_input), None)
+        weight_last_quantized_out = qsim2.model(dummy_input)
+
+        for enc1, enc2 in zip(qsim1.model.linear1.batch_matmul.param_quantizers['weight'].encoding,
+                              qsim2.model.linear1.batch_matmul.param_quantizers['weight'].encoding):
+            assert enc1.min == enc2.min
+            assert enc1.max == enc2.max
+            assert enc1.delta == enc2.delta
+            assert enc1.offset == enc2.offset
+            assert enc1.bw == enc2.bw
+
+        assert torch.equal(weight_first_quantized_out, weight_last_quantized_out)
+        assert len(qsim2.model.linear1.batch_matmul.param_quantizers['weight'].encoding) == 9
+
+        qsim2.export('./data', 'bmm_blockwise', dummy_input)
+
+class TestBlockwiseQuantGroupedConv:
+    def test_blockwise_quant_with_grouped_conv(self, create_per_channel_config):
+        torch.manual_seed(0)
+        dummy_input = torch.randn(1, 1, 9)
+        model = LinearModel()
+        orig_out = model(dummy_input)
+        orig_linear = model.linear1
+
+        model.linear1 = bqgc.LinearWithGroupedConv(orig_linear, 3)
+        new_out = model(dummy_input)
+
+        assert torch.allclose(orig_out, new_out, atol=1e-5)
+
+        qsim = QuantizationSimModel(model, dummy_input, config_file='./data/quantsim_config.json')
+        qsim.compute_encodings(lambda m, _: m(dummy_input), None)
+        _ = qsim.model(dummy_input)
+
+        assert len(qsim.model.linear1.grouped_conv.param_quantizers['weight'].encoding) == 9
+
+        qsim.export('./data', 'gc_blockwise', dummy_input)
+
+def test_parity_of_blockwise_quant_formulations(create_per_channel_config):
+
+    def assert_encodings_match(enc1, enc2):
+        assert enc1.min == enc2.min
+        assert enc1.max == enc2.max
+        assert enc1.delta == enc2.delta
+        assert enc1.offset == enc2.offset
+        assert enc1.bw == enc2.bw
+
+    torch.manual_seed(0)
+    dummy_input = torch.randn(1, 1, 9)
+    model_ts = LinearModel()
+    model_bm = copy.deepcopy(model_ts)
+    model_gc = copy.deepcopy(model_ts)
+
+    bqts.replace_linears_for_blockwise_quant(model_ts, 3)
+    bqbm.replace_linears_for_blockwise_quant(model_bm, 3)
+    bqgc.replace_linears_for_blockwise_quant(model_gc, 3)
+
+    qsim_ts = QuantizationSimModel(model_ts, dummy_input, config_file='./data/quantsim_config.json')
+    qsim_ts.compute_encodings(lambda m, _: m(dummy_input), None)
+
+    qsim_bm = QuantizationSimModel(model_bm, dummy_input, config_file='./data/quantsim_config.json')
+    bqbm.enable_bmm_per_channel_quantizers(qsim_bm)
+    qsim_bm.compute_encodings(lambda m, _: m(dummy_input), None)
+
+    qsim_gc = QuantizationSimModel(model_gc, dummy_input, config_file='./data/quantsim_config.json')
+    qsim_gc.compute_encodings(lambda m, _: m(dummy_input), None)
+
+    qsim_ts_linear1_encodings = list(itertools.chain(qsim_ts.model.linear1.linears[0].param_quantizers['weight'].encoding,
+                                                     qsim_ts.model.linear1.linears[1].param_quantizers['weight'].encoding,
+                                                     qsim_ts.model.linear1.linears[2].param_quantizers['weight'].encoding))
+
+    assert len(qsim_ts_linear1_encodings) == len(qsim_bm.model.linear1.batch_matmul.param_quantizers['weight'].encoding)
+    assert len(qsim_ts_linear1_encodings) == len(qsim_gc.model.linear1.grouped_conv.param_quantizers['weight'].encoding)
+    assert (len(qsim_ts.model.linear2.linears[0].param_quantizers['weight'].encoding) ==
+            len(qsim_bm.model.linear2.batch_matmul.param_quantizers['weight'].encoding))
+    assert (len(qsim_ts.model.linear2.linears[0].param_quantizers['weight'].encoding) ==
+            len(qsim_gc.model.linear2.grouped_conv.param_quantizers['weight'].encoding))
+
+    for idx, encoding in enumerate(qsim_ts_linear1_encodings):
+        assert_encodings_match(encoding, qsim_bm.model.linear1.batch_matmul.param_quantizers['weight'].encoding[idx])
+        assert_encodings_match(encoding, qsim_gc.model.linear1.grouped_conv.param_quantizers['weight'].encoding[idx])
+
+    for idx, encoding in enumerate(qsim_ts.model.linear2.linears[0].param_quantizers['weight'].encoding):
+        assert_encodings_match(encoding, qsim_bm.model.linear2.batch_matmul.param_quantizers['weight'].encoding[idx])
+        assert_encodings_match(encoding, qsim_gc.model.linear2.grouped_conv.param_quantizers['weight'].encoding[idx])
diff --git a/TrainingExtensions/torch/test/python/test_blockwise_quant_tensor_split.py b/TrainingExtensions/torch/test/python/test_blockwise_quant_tensor_split.py
deleted file mode 100644
index 318504bfcc..0000000000
--- a/TrainingExtensions/torch/test/python/test_blockwise_quant_tensor_split.py
+++ /dev/null
@@ -1,199 +0,0 @@
-# -*- mode: python -*-
-# =============================================================================
-#  @@-COPYRIGHT-START-@@
-#
-#  Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
-#
-#  Redistribution and use in source and binary forms, with or without
-#  modification, are permitted provided that the following conditions are met:
-#
-#  1. Redistributions of source code must retain the above copyright notice,
-#     this list of conditions and the following disclaimer.
-#
-#  2. Redistributions in binary form must reproduce the above copyright notice,
-#     this list of conditions and the following disclaimer in the documentation
-#     and/or other materials provided with the distribution.
-#
-#  3. Neither the name of the copyright holder nor the names of its contributors
-#     may be used to endorse or promote products derived from this software
-#     without specific prior written permission.
-#
-#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
-#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-#  POSSIBILITY OF SUCH DAMAGE.
-#
-#  SPDX-License-Identifier: BSD-3-Clause
-#
-#  @@-COPYRIGHT-END-@@
-# =============================================================================
-""" Tests for blockwise quant tensor split utility """
-
-import json
-import pytest
-import torch
-from aimet_torch.quantsim import QuantizationSimModel
-from aimet_torch.qc_quantize_op import QcQuantizeWrapper
-from aimet_torch.blockwise_quant_tensor_split import (BlockwiseLinear, replace_linears_for_blockwise_quant,
-                                                      tie_blockwise_linear_quantizers)
-
-class LinearModel(torch.nn.Module):
-    def __init__(self):
-        super(LinearModel, self).__init__()
-        self.linear1 = torch.nn.Linear(8, 3)
-        self.relu1 = torch.nn.ReLU()
-        self.linear2 = torch.nn.Linear(3, 2, bias=False)
-        self.softmax = torch.nn.Softmax()
-
-    def forward(self, inp):
-        x = self.linear1(inp)
-        x = self.relu1(x)
-        x = self.linear2(x)
-        x = self.softmax(x)
-        return x
-
-class InnerModel(torch.nn.Module):
-    def __init__(self):
-        super(InnerModel, self).__init__()
-        self.linear1 = torch.nn.Linear(8, 3)
-        self.seq = torch.nn.Sequential(torch.nn.Linear(3, 4), torch.nn.Linear(4, 5), torch.nn.ReLU())
-        self.seq2 = torch.nn.Sequential(torch.nn.Linear(5, 2))
-
-    def forward(self, inp):
-        x = self.linear1(inp)
-        x = self.seq(x)
-        x = self.seq2(x)
-        return x
-
-class NestedModel(torch.nn.Module):
-    def __init__(self):
-        super(NestedModel, self).__init__()
-        self.sigmoid = torch.nn.Sigmoid()
-        self.inner_linears = InnerModel()
-
-    def forward(self, inp):
-        x = self.inner_linears(inp)
-        x = self.sigmoid(x)
-        return x
-
-@pytest.fixture(scope='module')
-def create_per_channel_config():
-    quantsim_config = {
-        "defaults": {
-            "ops": {
-                "is_output_quantized": "True",
-                "is_symmetric": "False"
-            },
-            "params": {
-                "is_quantized": "False",
-                "is_symmetric": "True"
-            },
-            "per_channel_quantization": "True",
-        },
-        "params": {},
-        "op_type": {
-            'Split': {
-                'is_output_quantized': False
-            }
-        },
-        "supergroups": [],
-        "model_input": {},
-        "model_output": {}
-    }
-    with open('./data/quantsim_config.json', 'w') as f:
-        json.dump(quantsim_config, f)
-
-@pytest.mark.parametrize('model, dummy_input, block_size', [(torch.nn.Linear(8, 3), torch.randn(1, 8), 3),
-                                                             (torch.nn.Linear(8, 3, bias=False), torch.randn(1, 8), 3),
-                                                             (torch.nn.Linear(3, 2), torch.randn(1, 3), 3),
-                                                             (torch.nn.Linear(3, 2), torch.randn(1, 3), 4)])
-def test_blockwise_linear(model, dummy_input, block_size):
-    blockwise_linear = BlockwiseLinear(model, block_size=block_size)
-    orig_out = model(dummy_input)
-    new_out = blockwise_linear(dummy_input)
-    assert torch.allclose(orig_out, new_out, atol=1e-6)
-
-def test_replace_linears_for_blockwise_quant():
-    dummy_input = torch.randn(1, 8)
-    model = LinearModel()
-    linear1 = model.linear1
-    orig_out = model(dummy_input)
-    replace_linears_for_blockwise_quant(model, 3)
-
-    assert len(model.linear1.linears) == 3
-    assert torch.equal(model.linear1.linears[0].bias, linear1.bias)
-    for linear in model.linear1.linears[1:]:
-        assert linear.bias is None
-
-    assert len(model.linear1.elementwise_adds) == 2
-    assert len(model.linear2.linears) == 1
-    assert model.linear2.elementwise_adds is None
-    new_out = model(dummy_input)
-    assert torch.allclose(orig_out, new_out, atol=1e-6)
-
-def test_quantize_blockwise_linear(create_per_channel_config):
-    dummy_input = torch.randn(1, 8)
-    model = BlockwiseLinear(torch.nn.Linear(8, 3), 3)
-    qsim = QuantizationSimModel(model, dummy_input=dummy_input)
-
-    assert isinstance(qsim.model.split, QcQuantizeWrapper)
-    assert isinstance(qsim.model.linears[0], QcQuantizeWrapper)
-    assert isinstance(qsim.model.linears[1], QcQuantizeWrapper)
-    assert isinstance(qsim.model.linears[2], QcQuantizeWrapper)
-    assert isinstance(qsim.model.elementwise_adds[0], QcQuantizeWrapper)
-    assert isinstance(qsim.model.elementwise_adds[1], QcQuantizeWrapper)
-
-    # Temporary hack to disable split op output quantizers while handling for CG split op is reworked
-    for output_quantizer in qsim.model.split.output_quantizers:
-        output_quantizer.enabled = False
-
-    # Temporary hack to enable model input split op input quantizer while handling for CG split op is reworked
-    qsim.model.split.input_quantizers[0].enabled = True
-
-    qsim.compute_encodings(lambda m, _: m(dummy_input), None)
-    tie_blockwise_linear_quantizers(qsim)
-    _ = qsim.model(dummy_input)
-
-    assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
-            qsim.model.linears[1].output_quantizers[0].encoding.max)
-    assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
-            qsim.model.linears[2].output_quantizers[0].encoding.max)
-    assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
-            qsim.model.elementwise_adds[0].output_quantizers[0].encoding.max)
-    assert (qsim.model.linears[0].output_quantizers[0].encoding.max ==
-            qsim.model.elementwise_adds[1].output_quantizers[0].encoding.max)
-
-def test_blockwise_quant_with_small_linear(create_per_channel_config):
-    dummy_input = torch.randn(1, 3)
-    model = BlockwiseLinear(torch.nn.Linear(3, 2), 3)
-    qsim = QuantizationSimModel(model, dummy_input=dummy_input)
-    # Temporary hack to disable split op output quantizers while handling for CG split op is reworked
-    for output_quantizer in qsim.model.split.output_quantizers:
-        output_quantizer.enabled = False
-
-    # Temporary hack to enable model input split op input quantizer while handling for CG split op is reworked
-    qsim.model.split.input_quantizers[0].enabled = True
-
-    tie_blockwise_linear_quantizers(qsim)
-    qsim.compute_encodings(lambda m, _: m(dummy_input), None)
-    assert len(qsim.model.linears) == 1
-    assert qsim.model.elementwise_adds is None
-
-    _ = qsim.model(dummy_input)
-    assert len(qsim.connected_graph.get_all_ops()) == 1
-
-def test_nested_sequential_linears():
-    model = NestedModel()
-    replace_linears_for_blockwise_quant(model, block_size=2)
-    num_blockwise_linears = 0
-    for _, module in model.named_modules():
-        if isinstance(module, BlockwiseLinear):
-            num_blockwise_linears += 1
-    assert num_blockwise_linears == 4