nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/quantization/plugins/__init__.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Handles quantization plugins to correctly quantize third-party modules.
+
+Please check out the source code of this module for examples of how plugins work and how you can
+write your own one. Currently, we support plugins for
+
+- :meth:`apex<modelopt.torch.quantization.plugins.apex>`
+- :meth:`attention<modelopt.torch.quantization.plugins.attention>`
+- :meth:`diffusers<modelopt.torch.quantization.plugins.diffusers>`
+- :meth:`fairscale<modelopt.torch.quantization.plugins.fairscale>`
+- :meth:`huggingface<modelopt.torch.quantization.plugins.huggingface>`
+- :meth:`megatron<modelopt.torch.quantization.plugins.megatron>`
+- :meth:`peft<modelopt.torch.quantization.plugins.peft>`
+- :meth:`transformer_engine<modelopt.torch.quantization.plugins.transformer_engine>`
+"""
+
+from modelopt.torch.utils import import_plugin
+
+with import_plugin("accelerate"):
+    from .accelerate import *
+
+with import_plugin("apex"):
+    from .apex import *
+
+from .attention import *
+from .custom import *
+
+with import_plugin("diffusers"):
+    from .diffusers import *
+
+with import_plugin("fairscale"):
+    from .fairscale import *
+
+with import_plugin("huggingface"):
+    from .huggingface import *
+
+with import_plugin("megatron"):
+    from .megatron import *
+
+with import_plugin("nemo"):
+    from .nemo import *
+
+with import_plugin("peft"):
+    from .peft import *
+
+with import_plugin("transformer_engine"):
+    from .transformer_engine import *
+
+with import_plugin("transformers trainer"):
+    from .transformers_trainer import *
+
+with import_plugin("transformers"):
+    from .transformers import *
+
+with import_plugin("vllm"):
+    from .vllm import *
+
+with import_plugin("trl"):
+    from .trl import *

modelopt/torch/quantization/plugins/accelerate.py

@@ -0,0 +1,214 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Quantization support for accelerate modified models."""
+
+import warnings
+from contextlib import contextmanager
+from typing import Any
+
+import torch
+import transformers
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from accelerate.hooks import AlignDevicesHook, SequentialHook
+from accelerate.utils import get_max_memory, infer_auto_device_map
+from accelerate.utils.dataclasses import CustomDtype
+from accelerate.utils.offload import PrefixedDataset
+
+import modelopt.torch.quantization as mtq
+
+__all__ = ["init_quantized_weights"]
+
+
+def _get_cpu_offload_hook(hook):
+    if isinstance(hook, AlignDevicesHook) and hook.offload and hook.weights_map is not None:
+        assert "weight" in hook.weights_map
+        if (
+            isinstance(hook.weights_map, PrefixedDataset)
+            and hook.weights_map.prefix + "weight" not in hook.weights_map.dataset.state_dict
+        ):
+            raise NotImplementedError(
+                "This layer could be offloaded to disk. We don't support this yet."
+            )
+        return hook
+    elif isinstance(hook, SequentialHook):
+        for h in hook.hooks:
+            align_hook = _get_cpu_offload_hook(h)
+            if align_hook is not None:
+                return align_hook
+    return None
+
+
+@contextmanager
+def weight_access_and_writeback_context(module):
+    """Context manager for weight access and writeback for modules managed by accelerate."""
+    assert hasattr(module, "_hf_hook")
+    align_hook = _get_cpu_offload_hook(module._hf_hook)
+
+    if align_hook:
+        # Accelerate uses AlignDevicesHook to offload weights to CPU/Disk and then reload them in the forward pass
+        # The CPU/Disk offloaded weights are managed by PrefixDataset and OffloadedWeightsLoader
+        # See https://github.com/huggingface/accelerate/blame/f48d95c4939b281505a45b3d6e0bf554b65cc1ea/src/accelerate/utils/offload.py#L104-L141
+        # TODO: Add support for disk-offloaded models if needed (they will be really slow, hence low priority)
+
+        # This will load the weights from CPU state_dict and move it to the GPU from meta device
+        align_hook.pre_forward(module)
+    try:
+        yield
+    finally:
+        if align_hook:
+            # Update the weight in the CPU state_dict
+            if isinstance(align_hook.weights_map, PrefixedDataset):
+                key = align_hook.weights_map.prefix + "weight"
+                w_map = align_hook.weights_map.dataset.state_dict
+            else:
+                key, w_map = "weight", align_hook.weights_map
+            w_map[key] = module.weight.data.to(w_map[key].device, dtype=w_map[key].dtype)
+            align_hook.post_forward(module, None)
+
+
+@contextmanager
+def init_quantized_weights(
+    quant_cfg: dict[str, Any], gpu_mem_percentage: float = 0.8, quant_gemm: bool = False
+):
+    """Context manager for initializing and loading HuggingFace models with quantized and compressed weights.
+
+    This context manager patches `from_pretrained` to automatically:
+    1. Initialize the model with an empty config
+    2. Apply quantization configuration
+    3. Compress the weights
+    4. Load and dispatch the model to appropriate devices
+
+    Args:
+        quant_cfg: The quantization config to apply to the model
+        gpu_mem_percentage: Percentage of GPU memory to use (0.0-1.0)
+        quant_gemm: Whether to enable quantized GEMM
+
+    Example:
+        ```python
+        from transformers import AutoModelForCausalLM
+        from modelopt.torch.quantization.plugins.huggingface import init_quantized_weights
+        import modelopt.torch as mtq
+
+        # Load quantized and compressed model in one step
+        with init_quantized_weights(mtq.NVFP4_DEFAULT_CFG):
+            model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf")
+
+
+        # The model is already quantized, compressed, and loaded to appropriate devices
+
+        # calibrate model
+        mtq.calibrate(model, "max", forward_loop)
+        ```
+    """
+
+    def get_no_split_module_classes(model):
+        """Get no-split module classes for device mapping."""
+        no_split_classes = set()
+
+        # Get all named modules
+        for name, module in model.named_modules():
+            # Look for first layer patterns
+            if name.endswith((".layers.0", ".layer.0", ".h.0", ".blocks.0")):
+                no_split_classes.add(module.__class__.__name__)
+
+        return list(no_split_classes)
+
+    def get_model_device_map(model, gpu_mem_percentage):
+        """Create optimized device map for the quantized and compressed model."""
+        # Get compression ratio to adjust memory usage
+        max_memory = get_max_memory()
+        no_split_classes = get_no_split_module_classes(model)
+        if len(no_split_classes) == 0:
+            warnings.warn(
+                "No no-split module classes found for the model. Default device map might be incorrect."
+            )
+        max_memory = {k: v * gpu_mem_percentage for k, v in max_memory.items()}
+
+        def _get_special_dtype_map(model):
+            special_dtype_map = {}
+
+            def _get_byte_size(module):
+                weight_num_bit = module.weight_quantizer.num_bits
+                if isinstance(weight_num_bit, tuple):
+                    assert len(weight_num_bit) == 2, (
+                        "Weight num bit must be a tuple of two integers."
+                    )
+                    weight_num_bit = weight_num_bit[0] + weight_num_bit[1] + 1
+                if weight_num_bit == 8:
+                    return CustomDtype.FP8
+                elif weight_num_bit == 4:
+                    return CustomDtype.INT4
+                else:
+                    raise ValueError(f"Unsupported weight num bit: {weight_num_bit}")
+
+            for name, module in model.named_modules():
+                if (
+                    hasattr(module, "weight")
+                    and hasattr(module, "weight_quantizer")
+                    and not module.weight_quantizer.fake_quant
+                ):
+                    special_dtype_map[name + ".weight"] = _get_byte_size(module)
+            return special_dtype_map
+
+        special_dtype_map = _get_special_dtype_map(model)
+        inferred_device_map = infer_auto_device_map(
+            model,
+            max_memory=max_memory,
+            no_split_module_classes=no_split_classes,
+            special_dtypes=special_dtype_map,
+        )
+        return inferred_device_map
+
+    # Store original from_pretrained methods that we'll override
+    original_auto_causal_lm_from_pretrained = transformers.AutoModelForCausalLM.from_pretrained
+
+    # Create patched from_pretrained that applies quantization and compression
+    def patched_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs):
+        """Patched from_pretrained that handles quantization, compression and device mapping."""
+        # Initialize with empty config first to avoid loading weights twice
+        config = kwargs.pop("config", None)
+        trust_remote_code = kwargs.pop("trust_remote_code", False)
+        if config is None:
+            # Get config from pretrained model
+            from transformers import AutoConfig
+
+            config = AutoConfig.from_pretrained(
+                pretrained_model_name_or_path, trust_remote_code=trust_remote_code
+            )
+
+        with init_empty_weights():
+            # Fix torch_dtype to match original model
+            torch_dtype = kwargs.get("torch_dtype", getattr(config, "torch_dtype", torch.float16))
+            model = cls.from_config(config, torch_dtype=torch_dtype)
+
+        mtq.quantize(model, quant_cfg)
+        mtq.compress(model, config=mtq.CompressConfig(quant_gemm=quant_gemm))
+        _device_map = get_model_device_map(model, gpu_mem_percentage)
+
+        return load_checkpoint_and_dispatch(
+            model,
+            checkpoint=pretrained_model_name_or_path,
+            device_map=_device_map,
+            *args,
+            **kwargs,
+        )
+
+    try:
+        # Apply our patches
+        transformers.AutoModelForCausalLM.from_pretrained = classmethod(patched_from_pretrained)
+        yield
+    finally:
+        # Restore original methods
+        transformers.AutoModelForCausalLM.from_pretrained = original_auto_causal_lm_from_pretrained

modelopt/torch/quantization/plugins/apex.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Support quantization for apex linear layers."""
+
+from functools import partial
+
+import apex.transformer.tensor_parallel.layers as apex_parallel
+from apex.transformer.parallel_state import get_data_parallel_group, get_tensor_model_parallel_group
+
+from modelopt.torch.utils.distributed import ParallelState
+
+from ..nn import QuantModuleRegistry
+from ..nn.modules.quant_linear import _QuantLinear
+from .custom import _ParallelLinear
+
+
+class _ApexParallelLinear(_ParallelLinear):
+    def _setup(self):
+        quantized_linear_fn = partial(
+            _QuantLinear.quantized_linear_fn,
+            apex_parallel,
+            "linear_with_grad_accumulation_and_async_allreduce",
+            self,
+        )
+        self._forward_impl = quantized_linear_fn
+        self.parallel_state = ParallelState(
+            get_data_parallel_group(), get_tensor_model_parallel_group()
+        )
+        super()._setup()
+
+
+@QuantModuleRegistry.register({apex_parallel.ColumnParallelLinear: "apex_ColumnParallelLinear"})
+class _ApexColumnParallelLinear(_ApexParallelLinear):
+    _is_column_parallel = True
+
+
+@QuantModuleRegistry.register({apex_parallel.RowParallelLinear: "apex_RowParallelLinear"})
+class _ApexRowParallelLinear(_ApexParallelLinear):
+    _is_row_parallel = True

modelopt/torch/quantization/plugins/attention.py

@@ -0,0 +1,312 @@
+# Inspired by https://github.com/ELS-RD/transformer-deploy/blob/6b88e24ade6ce199e825adc0477b28a07f51f17d/src/transformer_deploy/QDQModels/ast_operator_patch.py
+
+# Apache License
+# Copyright 2022, Lefebvre Dalloz Services
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#    http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Support quantization for KV Cache in attention layers."""
+
+import ast
+import inspect
+import tempfile
+import types
+from warnings import warn
+
+from ..conversion import register
+from ..nn import TensorQuantizer
+
+__all__ = ["register_attention_for_kv_quant"]
+
+
+def register_attention_for_kv_quant(attention_cls: type) -> bool:
+    """Register attention layer for quantization of KV Cache.
+
+    Generate a quantized version of the attention class on the fly,
+    and register it with the original class for quantization.
+    """
+    source_code = inspect.getsource(attention_cls)
+    model_module = inspect.getmodule(attention_cls)
+    head = ast.parse(source_code)
+
+    bmm_ops = ("matmul", "bmm", "baddbmm")
+    sdpa_ops = ("scaled_dot_product_attention",)
+
+    def is_bmm(node):
+        return (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Attribute)
+            and node.func.attr in bmm_ops
+        )
+
+    def is_sdpa(node):
+        return (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Attribute)
+            and node.func.attr in sdpa_ops
+        )
+
+    def is_bin_matmul(node):
+        return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult)
+
+    def patch(node, quantizer_names, transpose=False):
+        for index, quantizer_name in enumerate(quantizer_names):
+            if quantizer_name is None:
+                continue
+            arg = node.args[index]
+
+            if not transpose:
+                node.args[index] = ast.Call(
+                    func=ast.Attribute(
+                        value=ast.Name(id="self", ctx=ast.Load()),
+                        attr=quantizer_name,
+                        ctx=ast.Load(),
+                    ),
+                    args=[arg],
+                    keywords=[],
+                )
+            else:
+                node.args[index] = ast.Call(
+                    func=ast.Attribute(
+                        value=ast.Call(
+                            func=ast.Attribute(
+                                value=ast.Name(id="self", ctx=ast.Load()),
+                                attr=quantizer_name,
+                                ctx=ast.Load(),
+                            ),
+                            args=[
+                                ast.Call(
+                                    func=ast.Attribute(value=arg, attr="transpose", ctx=ast.Load()),
+                                    args=[ast.Constant(value=-1), ast.Constant(value=-2)],
+                                    keywords=[],
+                                )
+                            ],
+                            keywords=[],
+                        ),
+                        attr="transpose",
+                        ctx=ast.Load(),
+                    ),
+                    args=[ast.Constant(value=-1), ast.Constant(value=-2)],
+                    keywords=[],
+                )
+
+    def patch_binop(node, quantizer_names, transpose=False):
+        assert len(quantizer_names) == 2
+        if quantizer_names[0] is not None:
+            node.left = ast.Call(
+                func=ast.Attribute(
+                    value=ast.Name(id="self", ctx=ast.Load()),
+                    attr=quantizer_names[0],
+                    ctx=ast.Load(),
+                ),
+                args=[node.left],
+                keywords=[],
+            )
+        if quantizer_names[1] is not None:
+            arg = node.right
+            if transpose:
+                arg = ast.Call(
+                    func=ast.Attribute(
+                        ast.Name(id="torch", ctx=ast.Load()),
+                        attr="transpose",
+                        ctx=ast.Load(),
+                    ),
+                    args=[arg, ast.Constant(value=-1), ast.Constant(value=-2)],
+                    keywords=[],
+                )
+            quant_arg = ast.Call(
+                func=ast.Attribute(
+                    value=ast.Name(id="self", ctx=ast.Load()),
+                    attr=quantizer_names[1],
+                    ctx=ast.Load(),
+                ),
+                args=[arg],
+                keywords=[],
+            )
+            if transpose:
+                quant_arg = ast.Call(
+                    func=ast.Attribute(
+                        ast.Name(id="torch", ctx=ast.Load()),
+                        attr="transpose",
+                        ctx=ast.Load(),
+                    ),
+                    args=[quant_arg, ast.Constant(value=-1), ast.Constant(value=-2)],
+                    keywords=[],
+                )
+            node.right = quant_arg
+
+    nodes = list(ast.walk(head))
+    org_class_name = nodes[1].name  # type: ignore[attr-defined]
+    new_class_name = nodes[1].name = "_Quant" + nodes[1].name  # type: ignore[attr-defined]
+
+    bmm_nodes = []
+    sdpa_nodes = []
+    bin_matmul_nodes = []
+    for node in ast.walk(head):
+        if is_bmm(node):
+            bmm_nodes.append(node)
+        if is_sdpa(node):
+            sdpa_nodes.append(node)
+        if is_bin_matmul(node):
+            bin_matmul_nodes.append(node)
+    if len(bmm_nodes) != 2 and len(sdpa_nodes) != 1 and len(bin_matmul_nodes) != 2:
+        print(f"Expect 2 bmm/matmul op in the {org_class_name}, found {len(bmm_nodes)}")
+        print(f"Or expect 1 sdpa op in the {org_class_name}, found {len(sdpa_nodes)}")
+        print(f"Or expect 2 @ op in the {org_class_name}, found {len(bin_matmul_nodes)}")
+        print("Auto quantization of KV Cache fails")
+        return False
+
+    if len(bmm_nodes) == 2:
+        # transpose k cache here to enable per-token quantization
+        # without transpose, the quantization will be per-channel, i.e.,
+        # self.k_bmm_quantizer(key_states.transpose(-1, -2))
+        # after transpose, the quantization will be per-token, i.e.,
+        # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2)
+        # removing the additional transpose is doable but not trivial
+        patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer"))
+        patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True)
+        print("Patching 2 BMM/Matmul operators with quantizers")
+    if len(bin_matmul_nodes) == 2:
+        patch_binop(
+            bin_matmul_nodes[1],
+            quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"),
+            transpose=True,
+        )
+        patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer"))
+        print("Patching 2 @ operators with quantizers")
+
+    if len(sdpa_nodes) == 1:
+        patch(
+            sdpa_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer")
+        )
+        print("Patching 1 scaled_dot_product_attention operator with quantizers")
+
+    head = ast.fix_missing_locations(head)
+    org_class = model_module.__dict__[org_class_name]
+
+    quant_class = _create_quantized_class_from_ast(head, org_class, new_class_name, model_module)
+    register(original_cls=org_class, quantized_cls=quant_class)
+    print(f"Successfully registered {org_class_name} for quantization")
+    return True
+
+
+def _create_quantized_class_from_ast(
+    head, org_class, new_class_name, model_module, temp_file_name=None
+):
+    """Create a quantized class from an AST representation.
+
+    Args:
+        head: The AST head containing the modified class definition
+        org_class: The original class to be quantized
+        new_class_name: Name for the new quantized class
+        model_module: The module containing the original class
+        temp_file_name: Optional file name to save the generated code
+
+    Returns:
+        The newly created quantized class
+    """
+    head = ast.fix_missing_locations(head)
+
+    # Save the generated code to a temporary file if requested
+    module_code_str = ast.unparse(head)
+    if temp_file_name is None:
+        with tempfile.NamedTemporaryFile(
+            prefix="modelopt_", suffix=".py", delete=False
+        ) as temp_file:
+            temp_file.write(module_code_str.encode())
+            temp_file_name = temp_file.name
+            print(f"Definition of {new_class_name} saved to {temp_file_name}")
+    else:
+        with open(temp_file_name, "w") as f:
+            f.write(module_code_str)
+
+    # Exec with python runtime and extract the new class
+    # This could lead to side effects if the class code is not properly isolated
+    # Therefore, it is recommended to run this function only when necessary
+    # exec(
+    #     new_class_code,
+    #     globals=model_module.__dict__,
+    #     locals=model_module.__dict__
+    # )  # bandit throws error here
+    # quant_class = model_module.__dict__[new_class_name]
+
+    # Extract the bytecode and create a new class on the fly
+    # This is more tricky but doesn't require runtime execution
+    module_code = compile(head, filename=f"{temp_file_name}", mode="exec")
+    class_code = module_code.co_consts[0]
+    assert class_code.co_name == new_class_name
+    method_codes = [const for const in class_code.co_consts if isinstance(const, types.CodeType)]
+
+    new_methods = {}
+    for method_code in method_codes:
+        method_name = method_code.co_name
+        original_method = getattr(org_class, method_name, None)
+        if not isinstance(original_method, types.FunctionType):
+            continue
+
+        # Check if the method is a decorated method
+        # The exec path can handle decorated methods, but the safety compliance disallows exec
+        closure = original_method.__closure__
+        globals = original_method.__globals__
+        if method_code.co_freevars != original_method.__code__.co_freevars:
+            warn(f"{new_class_name}.{method_name} is a decorated method. Ignoring the decorator!")
+
+            new_closure = ()
+            for freevar in method_code.co_freevars:
+                assert freevar in original_method.__closure__
+                new_closure += (
+                    original_method.__closure__[  # type: ignore[index]
+                        original_method.__code__.co_freevars.index(freevar)
+                    ],
+                )
+            closure = new_closure
+            for closure_item in original_method.__closure__:  # type: ignore[union-attr]
+                item = closure_item.cell_contents
+                if isinstance(item, types.FunctionType) and item.__name__ == method_name:
+                    globals = item.__globals__
+                    break
+            else:
+                raise ValueError(f"Cannot find the original method {method_name} in the closure")
+
+        # Create a new class method from bytecode
+        new_method = types.FunctionType(method_code, globals=globals, closure=closure)
+        new_method.__annotations__ = original_method.__annotations__
+        new_method.__defaults__ = original_method.__defaults__
+        new_method.__kwdefaults__ = original_method.__kwdefaults__
+        new_methods[method_name] = new_method
+
+    def setup_method(self):
+        self.q_bmm_quantizer = TensorQuantizer()
+        self.k_bmm_quantizer = TensorQuantizer()
+        self.v_bmm_quantizer = TensorQuantizer()
+
+    assert "_setup" not in new_methods, "Method _setup already exists"
+    new_methods["_setup"] = setup_method
+
+    # Create a new subclass on the fly
+    quant_class = type(new_class_name, (org_class,), new_methods)
+
+    return quant_class