transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/integrations/moe.py

@@ -0,0 +1,240 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from functools import wraps
+
+from ..utils.generic import GeneralInterface
+from ..utils.import_utils import is_torch_available
+
+
+if is_torch_available():
+    import torch
+
+
+# Examples of experts class with its eager mm implementation
+# class Experts(nn.Module):
+#     """Collection of expert weights stored as 3D tensors."""
+
+#     def __init__(self, config):
+#         super().__init__()
+#         self.num_experts = config.n_routed_experts
+#         self.hidden_dim = config.hidden_size
+#         self.intermediate_dim = config.moe_intermediate_size
+#         self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+#         self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+#         self.act_fn = ACT2FN[config.hidden_act]
+
+#     def forward(
+#         self,
+#         hidden_states: torch.Tensor,
+#         top_k_index: torch.Tensor,
+#         top_k_weights: torch.Tensor,
+#     ) -> torch.Tensor:
+#         final_hidden_states = torch.zeros_like(hidden_states)
+#         with torch.no_grad():
+#             expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+#             expert_mask = expert_mask.permute(2, 1, 0)
+#             expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+#         for expert_idx in expert_hit:
+#             expert_idx = expert_idx[0]
+#             if expert_idx == self.num_experts:
+#                 continue
+#             top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+#             current_state = hidden_states[token_idx]
+#             gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+#             current_hidden_states = self.act_fn(gate) * up
+#             current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+#             current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+#             final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+#         return final_hidden_states
+
+
+def batched_mm_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    device = hidden_states.device
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    num_experts = self.gate_up_proj.size(0)
+    final_hidden_states = torch.zeros_like(hidden_states)
+
+    # Flatten top_k_index to get expert_ids per selected sample
+    expert_ids = top_k_index.reshape(-1)
+    token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1)
+
+    # Resolve routing weights per selected sample, allowing top_k_weights to be either:
+    # - (num_tokens, num_top_k) Qwen2MoE style
+    # - (num_tokens, num_experts) DeepseekV2 style
+    if top_k_weights.shape == (num_tokens, num_top_k):
+        sample_weights = top_k_weights.reshape(-1)  # (S,)
+    elif top_k_weights.shape == (num_tokens, num_experts):
+        sample_weights = top_k_weights[token_idx, expert_ids]  # (S,)
+    else:
+        raise ValueError(
+            f"top_k_weights has an invalid/unsupported shape. It should be either (num_tokens, num_top_k)({num_tokens}, {num_top_k}) "
+            f"or (num_tokens, num_experts)({num_tokens}, {num_experts}), but got {top_k_weights.shape}."
+        )
+
+    # Get current hidden states for selected samples
+    current_hidden_states = hidden_states[token_idx]  # (S, hidden_dim)
+
+    # Select projection matrices for selected experts
+    selected_gate_up = self.gate_up_proj[expert_ids]  # (S, hidden_dim, 2 * intermediate_dim)
+    selected_down = self.down_proj[expert_ids]  # (S, hidden_dim, intermediate_dim)
+
+    # --- Up projection per expert (batched) ---
+    gate_up_out = torch.bmm(selected_gate_up, current_hidden_states.unsqueeze(-1)).squeeze(-1)
+    if hasattr(self, "gate_up_proj_bias") and self.gate_up_proj_bias is not None:
+        gate_up_out = gate_up_out + self.gate_up_proj_bias[expert_ids]
+
+    # Split into gate and up components
+    gate, up = gate_up_out.chunk(2, dim=-1)  # both have shape (S, intermediate_dim)
+
+    # Apply activation
+    hidden_after_activation = self.act_fn(gate) * up  # (S, intermediate_dim)
+
+    # --- Down projection per expert (batched) ---
+    out_per_sample = torch.bmm(selected_down, hidden_after_activation.unsqueeze(-1)).squeeze(-1)
+    if hasattr(self, "down_proj_bias") and self.down_proj_bias is not None:
+        out_per_sample = out_per_sample + self.down_proj_bias[expert_ids]
+
+    # Apply routing weights
+    out_per_sample = out_per_sample * sample_weights.unsqueeze(-1)  # (S, hidden_dim)
+
+    # Accumulate results back to the final_hidden_states using original token indices
+    final_hidden_states.index_add_(0, token_idx, out_per_sample.to(final_hidden_states.dtype))
+
+    return final_hidden_states
+
+
+def grouped_mm_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    if not hasattr(torch, "_grouped_mm"):
+        raise ImportError(
+            "torch._grouped_mm is not available. Please make sure you are using a PyTorch version that includes it (2.9+)."
+        )
+
+    device = hidden_states.device
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    num_experts = self.gate_up_proj.size(0)
+    final_hidden_states = torch.zeros_like(hidden_states)
+
+    # Flatten top_k_index to get expert_ids per selected sample
+    expert_ids = top_k_index.reshape(-1)
+    token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1)
+
+    # Get permutation to group by expert
+    perm = torch.argsort(expert_ids, stable=True)
+    inv_perm = torch.argsort(perm, stable=True)
+
+    # Resolve routing weights per selected sample, allowing top_k_weights to be either:
+    # - (num_tokens, num_top_k) Qwen2MoE style
+    # - (num_tokens, num_experts) DeepseekV2 style
+    if top_k_weights.shape == (num_tokens, num_top_k):
+        sample_weights = top_k_weights.reshape(-1)  # (S,)
+    elif top_k_weights.shape == (num_tokens, num_experts):
+        sample_weights = top_k_weights[token_idx, expert_ids]  # (S,)
+    else:
+        raise ValueError(
+            f"top_k_weights has an invalid/unsupported shape. It should be either (num_tokens, num_top_k)({num_tokens}, {num_top_k}) "
+            f"or (num_tokens, num_experts)({num_tokens}, {num_experts}), but got {top_k_weights.shape}."
+        )
+
+    # Get current hidden states for selected samples
+    current_hidden_states = hidden_states[token_idx]  # (S, hidden_dim)
+
+    # Group by expert for grouped_mm
+    expert_ids_g = expert_ids[perm]
+    sample_weights_g = sample_weights[perm]
+    current_states_g = current_hidden_states[perm]
+
+    # Compute offsets for grouped_mm
+    # using histc instead of bincount to avoid cuda graph issues
+    # (grouped_mm_experts_forward still fails with cuda graphs but because of _grouped_mm internals)
+    num_tokens_per_expert = torch.histc(expert_ids_g.float(), bins=num_experts, min=0, max=num_experts - 1)
+    offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)
+
+    # --- Up projection per expert (grouped_mm) ---
+    gate_up_out = torch._grouped_mm(current_states_g, self.gate_up_proj.transpose(-2, -1), offs=offsets)
+    if hasattr(self, "gate_up_proj_bias") and self.gate_up_proj_bias is not None:
+        # we should be able to pass bias to the grouped_mm call, but it's still not fully supported
+        gate_up_out = gate_up_out + self.gate_up_proj_bias[expert_ids_g]
+
+    # Split into gate and up components
+    gate, up = gate_up_out.chunk(2, dim=-1)  # both have shape (S, intermediate_dim)
+
+    # Apply activation
+    hidden_after_activation = self.act_fn(gate) * up  # (S, intermediate_dim)
+
+    # --- Down projection per expert (grouped_mm) ---
+    out_per_sample_g = torch._grouped_mm(hidden_after_activation, self.down_proj.transpose(-2, -1), offs=offsets)
+    if hasattr(self, "down_proj_bias") and self.down_proj_bias is not None:
+        # we should be able to pass bias to the grouped_mm call, but it's still not fully supported
+        out_per_sample_g = out_per_sample_g + self.down_proj_bias[expert_ids_g]
+
+    # Apply routing weights
+    out_per_sample_g = out_per_sample_g * sample_weights_g.unsqueeze(-1)
+
+    # Restore original order
+    out_per_sample = out_per_sample_g[inv_perm]
+
+    # Accumulate results back to the final_hidden_states using original token indices
+    final_hidden_states.index_add_(0, token_idx, out_per_sample.to(final_hidden_states.dtype))
+
+    return final_hidden_states
+
+
+class ExpertsInterface(GeneralInterface):
+    """Interface for registering custom experts implementations."""
+
+    _global_mapping = {
+        "batched_mm": batched_mm_experts_forward,
+        "grouped_mm": grouped_mm_experts_forward,
+    }
+
+
+ALL_EXPERTS_FUNCTIONS = ExpertsInterface()
+
+
+def use_experts_implementation(experts_class: type[torch.nn.Module]) -> type[torch.nn.Module]:
+    original_init = experts_class.__init__
+    original_forward = experts_class.forward
+
+    @wraps(original_init)
+    def __init__(self, config, *args, **kwargs):
+        original_init(self, config, *args, **kwargs)
+        self.config = config
+
+    @wraps(original_forward)
+    def forward(self, *args, **kwargs):
+        experts_forward = original_forward
+
+        if self.config._experts_implementation != "eager":
+            experts_forward = ALL_EXPERTS_FUNCTIONS[self.config._experts_implementation]
+
+        return experts_forward(self, *args, **kwargs)
+
+    experts_class.__init__ = __init__
+    experts_class.forward = forward
+    return experts_class

transformers/integrations/mxfp4.py

@@ -12,24 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging
+from ..utils import is_torch_available, is_torch_xpu_available, logging
 
 
 if is_torch_available():
     import torch
     from torch import nn
+from contextlib import contextmanager
 from typing import Optional
 
 from ..core_model_loading import ConversionOps
-
-
-if is_accelerate_available():
-    from accelerate import init_empty_weights
-
-import re
-from contextlib import contextmanager
-
-from ..quantizers.quantizers_utils import get_module_from_name
+from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module
 
 
 logger = logging.get_logger(__name__)
@@ -436,15 +429,6 @@ def mlp_forward(self, hidden_states):
     return routed_out, router_logits
 
 
-def should_convert_module(current_key_name, patterns):
-    current_key_name_str = ".".join(current_key_name)
-    if not any(
-        re.match(f"{key}\\.", current_key_name_str) or re.match(f"{key}", current_key_name_str) for key in patterns
-    ):
-        return True
-    return False
-
-
 def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs):
     from ..integrations.tensor_parallel import shard_and_distribute_module
 
@@ -604,70 +588,40 @@ def swizzle_mxfp4_convertops(blocks, scales, module, proj, target_device, triton
     )
 
 
-def _replace_with_mxfp4_linear(
-    model,
-    modules_to_not_convert=None,
-    current_key_name=None,
-    quantization_config=None,
-    has_been_replaced=False,
-    config=None,
-):
-    if current_key_name is None:
-        current_key_name = []
+def replace_with_mxfp4_linear(model, quantization_config=None, modules_to_not_convert: list[str] | None = None):
+    """
+    Public method that replaces the expert layers of the given model with mxfp4 quantized layers.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to convert, can be any `torch.nn.Module` instance.
+        quantization_config (`Mxfp4Config`, defaults to `None`):
+            The quantization config object that contains the quantization parameters.
+        modules_to_not_convert (`list`, *optional*, defaults to `None`):
+            A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be
+            converted.
+    """
+    if quantization_config.dequantize:
+        return model
+
+    from .hub_kernels import get_kernel
 
-    for name, module in model.named_children():
-        current_key_name.append(name)
-        if not should_convert_module(current_key_name, modules_to_not_convert):
-            current_key_name.pop(-1)
+    global triton_kernels_hub
+    triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
+
+    has_been_replaced = False
+    for module_name, module in model.named_modules():
+        if not should_convert_module(module_name, modules_to_not_convert):
             continue
         if module.__class__.__name__ == "GptOssExperts" and not quantization_config.dequantize:
-            with init_empty_weights():
-                model._modules[name] = Mxfp4GptOssExperts(config)
+            with torch.device("meta"):
+                model.set_submodule(module_name, Mxfp4GptOssExperts(model.config))
                 has_been_replaced = True
         if module.__class__.__name__ == "GptOssMLP" and not quantization_config.dequantize:
             from types import MethodType
 
             module.forward = MethodType(mlp_forward, module)
-        if len(list(module.children())) > 0:
-            _, has_been_replaced = _replace_with_mxfp4_linear(
-                module,
-                modules_to_not_convert,
-                current_key_name,
-                quantization_config,
-                has_been_replaced=has_been_replaced,
-                config=config,
-            )
-        current_key_name.pop(-1)
-    return model, has_been_replaced
-
 
-def replace_with_mxfp4_linear(
-    model,
-    modules_to_not_convert=None,
-    current_key_name=None,
-    quantization_config=None,
-    config=None,
-):
-    if quantization_config.dequantize:
-        return model
-    else:
-        from kernels import get_kernel
-
-        global triton_kernels_hub
-        triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
-
-    modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
-
-    if quantization_config.modules_to_not_convert is not None:
-        modules_to_not_convert.extend(quantization_config.modules_to_not_convert)
-    modules_to_not_convert = list(set(modules_to_not_convert))
-    model, has_been_replaced = _replace_with_mxfp4_linear(
-        model,
-        modules_to_not_convert,
-        current_key_name,
-        quantization_config,
-        config=config,
-    )
     if not has_been_replaced:
         logger.warning(
             "You are loading your model using mixed-precision FP4 quantization but no linear modules were found in your model."

transformers/integrations/peft.py

@@ -17,6 +17,7 @@ import json
 import os
 from typing import Any, Literal
 
+from ..conversion_mapping import get_model_conversion_mapping
 from ..core_model_loading import WeightRenaming, rename_source_key
 from ..utils import (
     CONFIG_NAME,
@@ -46,26 +47,6 @@ MIN_PEFT_VERSION = "0.18.0"
 logger = logging.get_logger(__name__)
 
 
-# DO NOT MODIFY, KEPT FOR BC ONLY
-VLMS = [
-    "aria",
-    "ayavision",
-    "emu3",
-    "fuyu",
-    "gotocr2",
-    "gemma3",
-    "internvl",
-    "llava",  # all llava prefixed models fall under this check
-    "mistral3",
-    "mllama",
-    "paligemma",
-    "qwen2vl",
-    "qwen2_5_vl",
-    "videollava",
-    "vipllava",
-]
-
-
 class PeftAdapterMixin:
     """
     A class containing all functions for loading and using adapters weights that are supported in PEFT library. For
@@ -103,6 +84,7 @@ class PeftAdapterMixin:
         low_cpu_mem_usage: bool = False,
         is_trainable: bool = False,
         hotswap: bool | Literal["auto"] = "auto",
+        local_files_only: bool = False,
         adapter_kwargs: dict[str, Any] | None = None,
     ) -> None:
         """
@@ -211,11 +193,10 @@ class PeftAdapterMixin:
             if any(conf.peft_type != PeftType.LORA for conf in self.peft_config.values()):
                 raise ValueError("Hotswapping is currently only supported for LoRA, please set `hotswap=False`.")
 
+        key_mapping = adapter_kwargs.pop("key_mapping", None) if adapter_kwargs is not None else None
+        weight_conversions = get_model_conversion_mapping(self, key_mapping=key_mapping)
         # peft only supports low_cpu_mem_usage starting from v0.13.0
         peft_load_kwargs = {}
-        key_mapping = adapter_kwargs.pop("key_mapping", None) if adapter_kwargs is not None else None
-        if key_mapping is None and any(allowed_name in self.__class__.__name__.lower() for allowed_name in VLMS):
-            key_mapping = self._checkpoint_conversion_mapping
         peft_load_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
 
         adapter_name = adapter_name if adapter_name is not None else "default"
@@ -263,6 +244,7 @@ class PeftAdapterMixin:
             adapter_config_file = find_adapter_config_file(
                 peft_model_id,
                 token=token,
+                local_files_only=local_files_only,
                 **adapter_kwargs,
             )
 
@@ -275,13 +257,11 @@ class PeftAdapterMixin:
             peft_config = PeftConfig.from_pretrained(
                 peft_model_id,
                 token=token,
+                local_files_only=local_files_only,
                 **adapter_kwargs,
             )
             peft_config.inference_mode = not is_trainable
 
-        if peft_config.peft_type != PeftType.LORA:
-            raise ValueError("Hotswapping is currently only supported for LoRA, please set `hotswap=False`.")
-
         if not hotswap:
             # TODO: WE NEED TOO APPLY OUR DYNAMIC WEIGHT CONVERSION AT SOME POINT HERE!
             # Create and add fresh new adapters into the model, unless the weights are hotswapped
@@ -291,21 +271,24 @@ class PeftAdapterMixin:
             self._hf_peft_config_loaded = True
 
         if peft_model_id is not None:
+            if "local_files_only" not in adapter_kwargs:
+                adapter_kwargs["local_files_only"] = local_files_only
             adapter_state_dict = load_peft_weights(peft_model_id, token=token, device=device, **adapter_kwargs)
 
         # We need to pre-process the state dict to remove unneeded prefixes - for backward compatibility
         renamings = []
-        if key_mapping:
-            renamings = [entry for entry in key_mapping if isinstance(entry, WeightRenaming)]
+        if weight_conversions:
+            renamings = [entry for entry in weight_conversions if isinstance(entry, WeightRenaming)]
         processed_adapter_state_dict = {}
         prefix = "base_model.model."
+        state_dict = self.state_dict()
         for key, value in adapter_state_dict.items():
             if key.startswith(prefix):
                 new_key = key[len(prefix) :]
             else:
                 new_key = key
 
-            new_key = rename_source_key(new_key, renamings, [])[0]
+            new_key = rename_source_key(new_key, renamings, [], self.base_model_prefix, state_dict)[0]
 
             # For hotswapping, we need the adapter name to be present in the state dict keys
             if hotswap:

transformers/integrations/quanto.py

@@ -12,21 +12,57 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..utils import is_optimum_quanto_available, is_torch_available, logging
+from ..core_model_loading import ConversionOps
+from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module
+from ..utils import is_torch_available, logging
 
 
 if is_torch_available():
     import torch
+    import torch.nn as nn
 
 logger = logging.get_logger(__name__)
 
 
+class QuantoQuantize(ConversionOps):
+    def __init__(self, hf_quantizer):
+        self.hf_quantizer = hf_quantizer
+
+    def convert(
+        self,
+        input_dict: dict[str, list[torch.Tensor]],
+        model: torch.nn.Module | None = None,
+        full_layer_name: str | None = None,
+        missing_keys: list[str] | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        _, value = tuple(input_dict.items())[0]
+        value = value[0]
+
+        from ..modeling_utils import _load_parameter_into_model
+
+        _load_parameter_into_model(model, full_layer_name, value)
+        module, _ = get_module_from_name(model, full_layer_name)
+        # Need to set those to a specific value, otherwise they will remain on meta device ...
+        module.input_scale = torch.ones(module.input_scale.shape)
+        module.output_scale = torch.ones(module.output_scale.shape)
+        # quantize
+        module.freeze()
+        module.weight.requires_grad = False
+        module._is_hf_initialized = True
+
+        # need to discard some missing keys we already updated the module in freeze.
+        module_name = full_layer_name.rsplit(".", 1)[0]
+        missing_keys.discard(f"{module_name}.weight")
+        missing_keys.discard(f"{module_name}.input_scale")
+        missing_keys.discard(f"{module_name}.output_scale")
+        return {}
+
+
 def replace_with_quanto_layers(
     model,
     quantization_config=None,
-    modules_to_not_convert=None,
-    current_key_name=None,
-    has_been_replaced=False,
+    modules_to_not_convert: list[str] | None = None,
 ):
     """
     Public method that recursively replaces the Linear layers of the given model with Quanto quantized layers.
@@ -35,64 +71,49 @@ def replace_with_quanto_layers(
     Args:
         model (`torch.nn.Module`):
             The model to convert, can be any `torch.nn.Module` instance.
-        quantization_config (`AqlmConfig`, defaults to `None`):
+        quantization_config (`QuantoConfig`, defaults to `None`):
             The quantization config object that contains the quantization parameters.
         modules_to_not_convert (`list`, *optional*, defaults to `None`):
             A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be
             converted.
-        current_key_name (`list`, *optional*, defaults to `None`):
-            A list that contains the current key name. This is used for recursion and should not be passed by the user.
-        has_been_replaced (`bool`, *optional*, defaults to `None`):
-            A boolean that indicates if the conversion has been successful or not. This is used for recursion and
-            should not be passed by the user.
     """
-    from accelerate import init_empty_weights
-
-    if is_optimum_quanto_available():
-        from optimum.quanto import QLayerNorm, QLinear, qfloat8, qint2, qint4, qint8
+    from optimum.quanto import QLayerNorm, QLinear, qfloat8, qint2, qint4, qint8
 
     w_mapping = {"float8": qfloat8, "int8": qint8, "int4": qint4, "int2": qint2}
     a_mapping = {None: None, "float8": qfloat8, "int8": qint8}
 
-    if modules_to_not_convert is None:
-        modules_to_not_convert = []
-
-    for name, module in model.named_children():
-        if current_key_name is None:
-            current_key_name = []
-        current_key_name.append(name)
-
-        if not any(key in ".".join(current_key_name) for key in modules_to_not_convert):
-            with init_empty_weights():
-                if isinstance(module, torch.nn.Linear):
-                    model._modules[name] = QLinear(
-                        in_features=module.in_features,
-                        out_features=module.out_features,
-                        bias=module.bias is not None,
-                        dtype=module.weight.dtype,
-                        weights=w_mapping[quantization_config.weights],
-                        activations=a_mapping[quantization_config.activations],
-                    )
-                    model._modules[name].requires_grad_(False)
-                    has_been_replaced = True
-                elif isinstance(module, torch.nn.LayerNorm):
-                    if quantization_config.activations is not None:
-                        model._modules[name] = QLayerNorm(
-                            module.normalized_shape,
-                            module.eps,
-                            module.elementwise_affine,
-                            module.bias is not None,
-                            activations=a_mapping[quantization_config.activations],
-                        )
-                        has_been_replaced = True
-        if len(list(module.children())) > 0:
-            _, has_been_replaced = replace_with_quanto_layers(
-                module,
-                quantization_config=quantization_config,
-                modules_to_not_convert=modules_to_not_convert,
-                current_key_name=current_key_name,
-                has_been_replaced=has_been_replaced,
-            )
-        # Remove the last key for recursion
-        current_key_name.pop(-1)
-    return model, has_been_replaced
+    has_been_replaced = False
+    for module_name, module in model.named_modules():
+        if not should_convert_module(module_name, modules_to_not_convert):
+            continue
+        with torch.device("meta"):
+            new_module = None
+            if isinstance(module, nn.Linear):
+                new_module = QLinear(
+                    in_features=module.in_features,
+                    out_features=module.out_features,
+                    bias=module.bias is not None,
+                    dtype=module.weight.dtype,
+                    weights=w_mapping[quantization_config.weights],
+                    activations=a_mapping[quantization_config.activations],
+                )
+            elif isinstance(module, torch.nn.LayerNorm) and quantization_config.activations is not None:
+                new_module = QLayerNorm(
+                    module.normalized_shape,
+                    module.eps,
+                    module.elementwise_affine,
+                    module.bias is not None,
+                    activations=a_mapping[quantization_config.activations],
+                )
+            if new_module is not None:
+                has_been_replaced = True
+                model.set_submodule(module_name, new_module)
+
+    if not has_been_replaced:
+        logger.warning(
+            "You are loading your model using quanto 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