[quantization] swap modules in one pass through a per-quantizer quantized_module (#50343)

* [quantization] swap modules in one pass through a per-quantizer quantized_module

* [quantization] build_fp8_linear / build_fp8_experts

Author: ArthurZucker

src/transformers/integrations/__init__.py

@@ -17,7 +17,6 @@
 
 
 _import_structure = {
-    "aqlm": ["replace_with_aqlm_linear"],
     "awq": [
         "post_init_awq_exllama_modules",
         "replace_quantization_scales",
@@ -26,7 +25,6 @@
     "bitnet": [
         "BitLinear",
         "pack_weights",
-        "replace_with_bitnet_linear",
         "unpack_weights",
     ],
     "bitsandbytes": [
@@ -48,9 +46,9 @@
         "set_hf_deepspeed_config",
         "unset_hf_deepspeed_config",
     ],
-    "eetq": ["replace_with_eetq_linear"],
+    "eetq": ["load_eetq_kernels"],
     "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"],
-    "finegrained_fp8": ["FP8Embedding", "FP8Linear", "replace_with_fp8_embedding", "replace_with_fp8_linear"],
+    "finegrained_fp8": ["FP8Embedding", "FP8Linear"],
     "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"],
     "gemma_quant": [
         "QuantizedEmbedding",
@@ -67,7 +65,6 @@
         "HiggsLinear",
         "dequantize_higgs",
         "quantize_with_higgs",
-        "replace_with_higgs_linear",
     ],
     "hqq": ["prepare_for_hqq_linear"],
     "hub_kernels": [
@@ -127,7 +124,6 @@
     "liger": ["apply_liger_kernel"],
     "metal_quantization": [
         "MetalLinear",
-        "replace_with_metal_linear",
     ],
     "moe": [
         "batched_mm_experts_forward",
@@ -146,7 +142,7 @@
         "deactivate_neftune",
         "neftune_post_forward_hook",
     ],
-    "nvfp4": ["NVFP4Linear", "NVFP4Quantize", "replace_with_nvfp4_linear"],
+    "nvfp4": ["NVFP4Linear", "NVFP4Quantize"],
     "peft": ["PeftAdapterMixin"],
     "quanto": ["replace_with_quanto_layers"],
     "sinq": ["SinqDeserialize", "SinqQuantize"],
@@ -174,7 +170,6 @@
 ]
 
 if TYPE_CHECKING:
-    from .aqlm import replace_with_aqlm_linear
     from .awq import (
         post_init_awq_exllama_modules,
         replace_quantization_scales,
@@ -183,7 +178,6 @@
     from .bitnet import (
         BitLinear,
         pack_weights,
-        replace_with_bitnet_linear,
         unpack_weights,
     )
     from .bitsandbytes import (
@@ -205,9 +199,9 @@
         set_hf_deepspeed_config,
         unset_hf_deepspeed_config,
     )
-    from .eetq import replace_with_eetq_linear
+    from .eetq import load_eetq_kernels
     from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear
-    from .finegrained_fp8 import FP8Embedding, FP8Linear, replace_with_fp8_embedding, replace_with_fp8_linear
+    from .finegrained_fp8 import FP8Embedding, FP8Linear
     from .fsdp import is_fsdp_enabled, is_fsdp_managed_module
     from .gemma_quant import (
         QuantizedEmbedding,
@@ -220,7 +214,7 @@
         GGUF_CONFIG_MAPPING,
         _gguf_parse_value,
     )
-    from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear
+    from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs
     from .hqq import prepare_for_hqq_linear
     from .hub_kernels import (
         LayerRepository,
@@ -279,7 +273,6 @@
     from .liger import apply_liger_kernel
     from .metal_quantization import (
         MetalLinear,
-        replace_with_metal_linear,
     )
     from .moe import (
         batched_mm_experts_forward,
@@ -293,7 +286,7 @@
         swizzle_mxfp4,
     )
     from .neftune import activate_neftune, deactivate_neftune, neftune_post_forward_hook
-    from .nvfp4 import NVFP4Linear, NVFP4Quantize, replace_with_nvfp4_linear
+    from .nvfp4 import NVFP4Linear, NVFP4Quantize
     from .peft import PeftAdapterMixin
     from .quanto import replace_with_quanto_layers
     from .sinq import SinqDeserialize, SinqQuantize

src/transformers/integrations/aqlm.py

@@ -1,5 +1,5 @@
 from ..core_model_loading import ConversionOps
-from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module
+from ..quantizers.quantizers_utils import get_module_from_name
 from ..utils import is_torch_available, logging
 
 
@@ -282,64 +282,6 @@ def forward(self, input):
         return output
 
 
-def replace_with_bitnet_linear(model, modules_to_not_convert: list[str] | None = None, quantization_config=None):
-    """
-    Public method that replaces the linear layers of the given model with bitnet quantized layers.
-
-    Args:
-        model (`torch.nn.Module`):
-            The model to convert, can be any `torch.nn.Module` instance.
-        modules_to_not_convert (`list[str]`, *optional*, defaults to `None`):
-            A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be
-            converted.
-        quantization_config (`BitNetConfig`):
-            The quantization config object that contains the quantization parameters.
-    """
-
-    has_been_replaced = False
-    # we need this to correctly materialize the weights during quantization
-    for module_name, module in model.named_modules():
-        if not should_convert_module(module_name, modules_to_not_convert):
-            continue
-        with torch.device("meta"):
-            if isinstance(module, nn.Linear):
-                if quantization_config and quantization_config.linear_class == "autobitlinear":
-                    new_module = AutoBitLinear(
-                        in_features=module.in_features,
-                        out_features=module.out_features,
-                        bias=module.bias is not None,
-                        device=module.weight.device,
-                        dtype=module.weight.dtype,
-                        online_quant=(quantization_config.quantization_mode == "online"),
-                        use_rms_norm=quantization_config.use_rms_norm,
-                        rms_norm_eps=quantization_config.rms_norm_eps,
-                    )
-                    if quantization_config.quantization_mode == "offline":
-                        new_module.requires_grad_(False)
-                else:
-                    new_module = BitLinear(
-                        in_features=module.in_features,
-                        out_features=module.out_features,
-                        bias=module.bias is not None,
-                        device=module.weight.device,
-                        dtype=module.weight.dtype,
-                        use_rms_norm=quantization_config.use_rms_norm if quantization_config else False,
-                        rms_norm_eps=quantization_config.rms_norm_eps if quantization_config else 1e-6,
-                    )
-                    new_module.requires_grad_(False)
-                model.set_submodule(module_name, new_module)
-                has_been_replaced = True
-
-    if not has_been_replaced:
-        logger.warning(
-            "You are loading your model using bitnet but no linear modules were found in your model."
-            " Please double check your model architecture, or submit an issue on github if you think this is"
-            " a bug."
-        )
-
-    return model
-
-
 class BitNetDeserialize(ConversionOps):
     def __init__(self, hf_quantizer):
         self.hf_quantizer = hf_quantizer

src/transformers/integrations/bitnet.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from ..core_model_loading import ConversionOps
-from ..quantizers.quantizers_utils import should_convert_module
 from ..utils import is_torch_available, logging
 
 
@@ -83,41 +82,9 @@ def forward(self, input):
         return output
 
 
-def replace_with_eetq_linear(model, modules_to_not_convert: list[str] | None = None, pre_quantized=False):
-    """
-    A helper function to replace all `torch.nn.Linear` modules by `EetqLinear` modules.
-
-    Parameters:
-        model (`torch.nn.Module`):
-            Input model or `torch.nn.Module` as the function is run recursively.
-        modules_to_not_convert (`list[`str`]`, *optional*, defaults to `None`):
-            Names of the modules to not convert in `EetqLinear`. In practice we keep the `lm_head` in full precision
-            for numerical stability reasons.
-    """
+def load_eetq_kernels():
+    """Loads the EETQ hub kernels used by `EetqLinear` and `EetqQuantize`."""
     from .hub_kernels import get_kernel
 
     global eetq_kernels_hub
     eetq_kernels_hub = get_kernel("kernels-community/quantization-eetq", version=1)
-
-    has_been_replaced = False
-    # we need this to correctly materialize the weights during quantization
-    module_kwargs = {} if pre_quantized else {"dtype": None}
-    for module_name, module in model.named_modules():
-        if not should_convert_module(module_name, modules_to_not_convert):
-            continue
-        with torch.device("meta"):
-            if isinstance(module, nn.Linear):
-                new_module = EetqLinear(
-                    module.in_features, module.out_features, bias=module.bias is not None, **module_kwargs
-                )
-                model.set_submodule(module_name, new_module)
-                has_been_replaced = True
-
-    if not has_been_replaced:
-        logger.warning(
-            "You are loading your model using eetq but no linear modules were found in your model."
-            " Please double check your model architecture, or submit an issue on github if you think this is"
-            " a bug."
-        )
-
-    return model