Repository: huggingface/transformers
Pull request: [device_map] Accelerate loading by computing device_map much faster

# What does this PR do?

If a model has a lot of different parameters (which is the case for models with a lot of experts, such as Qwen3Next), the computation of the device_map itself (yes, just the mapping of modules to devices) is EXTREMELY slow when using `accelerate`.
This is because they traverse the model graph in a very non-optimized manner, resulting in quadratic complexity in the number of parameters and modules, when we could do it in linear time instead. For example, `Qwen/Qwen3-Next-80B-A3B-Instruct` has a total of `173529` parameters and intermediate modules, which is a a lot. So doing the following:

```python
from transformers import AutoConfig, Qwen3NextForCausalLM
from transformers.modeling_utils import _get_device_map
import torch
import time

model_name = "Qwen/Qwen3-Next-80B-A3B-Instruct"
config = AutoConfig.from_pretrained(model_name)

with torch.device("meta"):
    model = Qwen3NextForCausalLM(config)

t0 = time.time()
device_map = _get_device_map(model, device_map="auto", dtype=torch.bfloat16, max_memory=None, hf_quantizer=None, keep_in_fp32_regex=None)
dt = time.time() - t0
print(f"Time needed: {dt:.2f} s")

>>> Time needed: 60.02 s
```

takes un UNBEARABLE `60.02` s on my Macbook Pro M3. And this is only for the device_map computation, the loading has not even started yet...

This PR fixes the issue by rewriting the needed functions from `accelerate` with much better complexity and operations directly in Transformers, leading to linear time and the following:

```python
>>> Time needed: 1.33 s
```

So effectively going 60x faster from 1min to 1s

The device_map created are the exact same.

Changed files:
filename: src/transformers/integrations/accelerate.py
patch: @@ -12,20 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Since, https://github.com/huggingface/transformers/pull/36963, loading is always performed with models on meta
-device. But since the `init_empty_weights` and `find_tied_parameters` functions are from accelerate, and accelerate is
-somewhat still a soft dependency, we copy the functions here to be used natively in Transformers.
-
-The `init_empty_weights` and `init_on_device` functions were copied from `accelerate.big_modeling.py`, and the
-`find_tied_parameters` was copied from `accelerate.utils.modeling.py`
+Some of the functions here are derived from the `accelerate` library, with some tweaks for better performances
+and simplicity/ease of use.
 """
 
-import collections
+import copy
 import inspect
 import os
+from collections import defaultdict
 from contextlib import contextmanager
-
-from ..utils import is_accelerate_available, is_torch_available, logging
+from typing import TYPE_CHECKING
+
+from ..utils import (
+    is_accelerate_available,
+    is_torch_available,
+    is_torch_xpu_available,
+    logging,
+)
 from ..utils.quantization_config import QuantizationMethod
 from .deepspeed import is_deepspeed_zero3_enabled
 from .fsdp import is_fsdp_enabled
@@ -36,7 +39,12 @@
     import torch.nn as nn
 
 if is_accelerate_available():
-    from accelerate import dispatch_model
+    from accelerate import dispatch_model, infer_auto_device_map
+    from accelerate.utils import check_tied_parameters_on_same_device, get_max_memory
+
+if TYPE_CHECKING:
+    from ..modeling_utils import PreTrainedModel
+    from ..quantizers import HfQuantizer
 
 
 logger = logging.get_logger(__name__)
@@ -205,7 +213,7 @@ def find_tied_parameters(model: "nn.Module", **kwargs):
     return [sorted([weight] + list(set(tied))) for weight, tied in tied_param_groups.items()]
 
 
-def check_and_set_device_map(device_map):
+def check_and_set_device_map(device_map: "torch.device | int | str | dict | None") -> dict | str | None:
     from ..modeling_utils import get_torch_context_manager_or_global_device
 
     # Potentially detect context manager or global device, and use it (only if no device_map was provided)
@@ -249,6 +257,203 @@ def check_and_set_device_map(device_map):
     return device_map
 
 
+def compute_module_sizes(
+    model: "PreTrainedModel", hf_quantizer: "HfQuantizer | None"
+) -> tuple[dict[str, int], dict[str, int]]:
+    """
+    Compute the size of each submodule of a given model (in bytes).
+    Returns a tuple of 2 dicts, the fist one containing a mapping of all the modules and the corresponding size
+    in bytes, and the 2nd one containing a mapping from all leaf modules (modules containing parameters, the end of
+    the model graph) and the corresponding sizes.
+    """
+    all_module_sizes = defaultdict(int)
+    leaves_module_sizes = defaultdict(int)
+    for name, param in model.state_dict().items():
+        if hf_quantizer is not None:
+            dtype_size = hf_quantizer.param_element_size(model, name)
+        else:
+            dtype_size = param.element_size()
+        size = param.numel() * dtype_size
+        name_parts = name.split(".")
+        for idx in range(len(name_parts)):
+            all_module_sizes[".".join(name_parts[:idx])] += size
+        if "." in name:
+            leaves_module_sizes[name.rsplit(".", 1)[0]] += size
+
+    return all_module_sizes, leaves_module_sizes
+
+
+def get_balanced_memory(
+    model: "PreTrainedModel",
+    max_memory: dict[int | str, int | str] | None = None,
+    no_split_module_classes: list[str] | None = None,
+    hf_quantizer: "HfQuantizer | None" = None,
+    low_zero: bool = False,
+):
+    """
+    Compute a `max_memory` dictionary for [`infer_auto_device_map`] that will balance the use of each available GPU.
+
+    <Tip>
+
+    All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
+    meta device (as it would if initialized within the `init_empty_weights` context manager).
+
+    </Tip>
+
+    Args:
+        model (`PreTrainedModel`):
+            The model to analyze.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
+            Example: `max_memory={0: "1GB"}`.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        hf_quantizer (`HfQuantizer`, *optional*):
+            A quantizer for the model.
+        low_zero (`bool`, *optional*):
+            Minimizes the number of weights on GPU 0, which is convenient when it's used for other operations (like the
+            Transformers generate function).
+    """
+    # Get default / clean up max_memory
+    user_not_set_max_memory = max_memory is None
+    max_memory = get_max_memory(max_memory)
+    # Check the number of accelerators available
+    accelerator_max_memory = copy.deepcopy(max_memory)
+    _, _ = accelerator_max_memory.pop("cpu", None), accelerator_max_memory.pop("disk", None)
+    num_devices = len([d for d in accelerator_max_memory if accelerator_max_memory[d] > 0])
+
+    if num_devices == 0:
+        return max_memory
+
+    if num_devices == 1:
+        # We cannot do low_zero on just one GPU, but we will still reserve some memory for the buffer
+        low_zero = False
+        # If user just asked us to handle memory usage, we should avoid OOM
+        if user_not_set_max_memory:
+            for key in max_memory.keys():
+                if isinstance(key, int):
+                    max_memory[key] *= 0.9  # 90% is a good compromise
+                    logger.info(
+                        f"We will use 90% of the memory on device {key} for storing the model, and 10% for the buffer to avoid OOM. "
+                        "You can set `max_memory` in to a higher value to use more memory (at your own risk)."
+                    )
+                    break  # only one device
+
+    module_sizes, leave_modules_sizes = compute_module_sizes(model, hf_quantizer)
+    per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
+
+    # We can't just set the memory to model_size // num_devices as it will end being too small: each GPU will get
+    # slightly less layers and some layers will end up offload at the end. So this function computes a buffer size to
+    # add which is the biggest of:
+    # - the size of no split block (if applicable)
+    # - the mean of the layer sizes
+    if no_split_module_classes is None:
+        no_split_module_classes = []
+    elif not isinstance(no_split_module_classes, (list, tuple)):
+        no_split_module_classes = [no_split_module_classes]
+
+    # Identify the size of the no_split_block modules
+    buffer = 0
+    if len(no_split_module_classes) > 0:
+        no_split_children = {}
+        for name, size in module_sizes.items():
+            if name == "":
+                continue
+            submodule = model.get_submodule(name)
+            class_name = submodule.__class__.__name__
+            if class_name in no_split_module_classes and class_name not in no_split_children:
+                no_split_children[class_name] = size
+
+            if set(no_split_children.keys()) == set(no_split_module_classes):
+                break
+        buffer = max(no_split_children.values()) if len(no_split_children) > 0 else 0
+
+    mean_leaves = int(sum(leave_modules_sizes.values()) / max(len(leave_modules_sizes), 1))
+    buffer = int(1.25 * max(buffer, mean_leaves))
+    per_gpu += buffer
+
+    # Sorted list of GPUs id (we may have some gpu ids not included in the our max_memory list - let's ignore them)
+    gpus_idx_list = sorted(
+        device_id for device_id, device_mem in max_memory.items() if isinstance(device_id, int) and device_mem > 0
+    )
+    # The last device is left with max_memory just in case the buffer is not enough.
+    for idx in gpus_idx_list[:-1]:
+        max_memory[idx] = min(max_memory[0] if low_zero and idx == 0 else per_gpu, max_memory[idx])
+
+    if low_zero:
+        min_zero = max(0, module_sizes[""] - sum([max_memory[i] for i in range(1, num_devices)]))
+        max_memory[0] = min(min_zero, max_memory[0])
+
+    return max_memory
+
+
+def _get_device_map(
+    model: "PreTrainedModel",
+    device_map: dict | str | None,
+    max_memory: dict | None,
+    hf_quantizer: "HfQuantizer | None",
+    dtype: torch.dtype | None,
+) -> dict:
+    """Compute the final `device_map` to use if we passed a value in ['auto', 'balanced', 'balanced_low_0', 'sequential'].
+    Otherwise, we check for any device inconsistencies in the device_map.
+    """
+    if isinstance(device_map, str):
+        special_dtypes = {}
+        if hf_quantizer is not None:
+            special_dtypes = hf_quantizer.get_special_dtypes_update(model, dtype)
+
+        target_dtype = dtype
+        if hf_quantizer is not None:
+            target_dtype = hf_quantizer.adjust_target_dtype(target_dtype)
+
+        no_split_modules = model._get_no_split_modules(device_map)
+
+        if device_map != "sequential":
+            inferred_max_memory = get_balanced_memory(
+                model,
+                max_memory=max_memory,
+                no_split_module_classes=no_split_modules,
+                hf_quantizer=hf_quantizer,
+                low_zero=(device_map == "balanced_low_0"),
+            )
+        else:
+            inferred_max_memory = get_max_memory(max_memory)
+        if hf_quantizer is not None:
+            inferred_max_memory = hf_quantizer.adjust_max_memory(inferred_max_memory)
+
+        # `inferred_max_memory` contains non-reserved memory. There may be *unused* reserved memory in the GPU,
+        # which we can use to allocate parameters.
+        for device_name in inferred_max_memory:
+            if isinstance(device_name, int):  # it's a GPU device
+                if is_torch_xpu_available():
+                    unused_memory = torch.xpu.memory_reserved(device_name) - torch.xpu.memory_allocated(device_name)
+                else:
+                    unused_memory = torch.cuda.memory_reserved(device_name) - torch.cuda.memory_allocated(device_name)
+                inferred_max_memory[device_name] += unused_memory
+            # respect the `max_memory` passed by the user
+            if max_memory is not None and device_name in max_memory:
+                inferred_max_memory[device_name] = min(inferred_max_memory[device_name], max_memory[device_name])
+
+        device_map = infer_auto_device_map(
+            model,
+            max_memory=inferred_max_memory,
+            dtype=target_dtype,
+            no_split_module_classes=no_split_modules,
+            special_dtypes=special_dtypes,
+        )
+
+        if hf_quantizer is not None:
+            hf_quantizer.validate_environment(device_map=device_map)
+
+    elif device_map is not None:
+        tied_params = find_tied_parameters(model)
+        # check if we don't have tied param in different devices
+        check_tied_parameters_on_same_device(tied_params, device_map)
+
+    return device_map
+
+
 def accelerate_dispatch(model, hf_quantizer, device_map, offload_folder, offload_index, offload_buffers):
     device_map_kwargs = {
         "device_map": device_map,
@@ -281,7 +486,7 @@ def get_disk_only_shard_files(device_map, weight_map):
     """
     Returns the list of shard files containing only weights offloaded to disk.
     """
-    files_content = collections.defaultdict(list)
+    files_content = defaultdict(list)
     for weight_name, filename in weight_map.items():
         while len(weight_name) > 0 and weight_name not in device_map:
             weight_name = ".".join(weight_name.split(".")[:-1])

filename: src/transformers/modeling_utils.py
patch: @@ -50,6 +50,7 @@
 from .generation import CompileConfig, GenerationConfig
 from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled, is_fsdp_enabled
 from .integrations.accelerate import (
+    _get_device_map,
     accelerate_disk_offload,
     accelerate_dispatch,
     check_and_set_device_map,
@@ -109,7 +110,6 @@
     is_torch_mlu_available,
     is_torch_npu_available,
     is_torch_xla_available,
-    is_torch_xpu_available,
     logging,
 )
 from .utils.generic import _CAN_RECORD_REGISTRY, GeneralInterface, OutputRecorder
@@ -125,13 +125,9 @@
 
 
 if is_accelerate_available():
-    from accelerate import infer_auto_device_map
     from accelerate.hooks import add_hook_to_module
     from accelerate.utils import (
-        check_tied_parameters_on_same_device,
         extract_model_from_parallel,
-        get_balanced_memory,
-        get_max_memory,
         offload_weight,
         save_offload_index,
     )
@@ -277,27 +273,6 @@ def _wrapper(*args, **kwargs):
     return _wrapper
 
 
-def get_keep_in_fp32_regex(model, hf_quantizer, dtype):
-    # Find fp32 modules if needed
-    keep_in_fp32_modules = []
-    # The _keep_in_fp32_modules flag is only used to avoid bf16 -> fp16 casting precision issues. It was introduced
-    # in case of force loading a model that should stay bf16 in fp16 (which includes a few quantizers as this is a pre-processing
-    # step for e.g. bitsandbytes). See https://github.com/huggingface/transformers/issues/20287 for details.
-    if model._keep_in_fp32_modules is not None and (
-        dtype == torch.float16 or getattr(hf_quantizer, "use_keep_in_fp32_modules", False)
-    ):
-        keep_in_fp32_modules.extend(model._keep_in_fp32_modules)
-
-    if model._keep_in_fp32_modules_strict is not None and (dtype == torch.float16 or dtype == torch.bfloat16):
-        keep_in_fp32_modules.extend(model._keep_in_fp32_modules_strict)
-
-    keep_in_fp32_regex = None
-    if keep_in_fp32_modules:
-        # We need to match exact layers, so we add either `.` on each side, or start/end of string
-        keep_in_fp32_regex = re.compile("|".join([rf"((^|\.){module}($|\.))" for module in keep_in_fp32_modules]))
-    return keep_in_fp32_regex
-
-
 def get_torch_context_manager_or_global_device():
     """
     Test if a device context manager is currently in use, or if it is not the case, check if the default device
@@ -559,7 +534,6 @@ def _infer_parameter_dtype(
     model: "PreTrainedModel",
     param_name: str,
     empty_param: torch.Tensor,
-    keep_in_fp32_regex: Optional[re.Pattern] = None,
     hf_quantizer: Optional[HfQuantizer] = None,
 ) -> Union[bool, Optional[torch.dtype]]:
     try:
@@ -580,11 +554,8 @@ def _infer_parameter_dtype(
     casting_dtype = None
     is_param_float8_e4m3fn = is_torch_e4m3fn_available and empty_param.dtype == torch.float8_e4m3fn
     if empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn:
-        # First fp32 if part of the exception list
-        if keep_in_fp32_regex is not None and keep_in_fp32_regex.search(param_name):
-            casting_dtype = torch.float32
-        # Then dtype that was instantiated in the meta model -- note that this respects subconfigs dtypes
-        elif hf_quantizer is not None:
+        # dtype that was instantiated in the meta model -- note that this respects subconfigs dtypes
+        if hf_quantizer is not None:
             casting_dtype = model.config._pre_quantization_dtype
         else:
             casting_dtype = old_param.dtype
@@ -608,7 +579,6 @@ def _load_state_dict_into_meta_model(
     disk_offload_folder: Optional[str] = None,
     disk_offload_index: Optional[dict] = None,
     hf_quantizer: Optional[HfQuantizer] = None,
-    keep_in_fp32_regex: Optional[re.Pattern] = None,
     device_mesh: Optional["torch.distributed.device_mesh.DeviceMesh"] = None,
 ) -> tuple[Optional[dict], Optional[dict]]:
     """Load parameters from `meta_state_dict` into the model. The parameters of the `meta_state_dict` are on the meta
@@ -638,13 +608,7 @@ def _load_state_dict_into_meta_model(
             param = file_pointer.get_slice(serialized_param_name)
         else:
             param = empty_param.to(tensor_device)  # It is actually not empty!
-        to_contiguous, casting_dtype = _infer_parameter_dtype(
-            model,
-            param_name,
-            empty_param,
-            keep_in_fp32_regex,
-            hf_quantizer,
-        )
+        to_contiguous, casting_dtype = _infer_parameter_dtype(model, param_name, empty_param, hf_quantizer)
 
         if device_mesh is not None:
             if not is_quantized or not hf_quantizer.param_needs_quantization(model, param_name):
@@ -745,7 +709,6 @@ def load_shard_file(args):
         reverse_key_renaming_mapping,
         disk_offload_folder,
         disk_offload_index,
-        keep_in_fp32_regex,
         device_mesh,
     ) = args
 
@@ -780,7 +743,6 @@ def load_shard_file(args):
             disk_offload_folder=disk_offload_folder,
             disk_offload_index=disk_offload_index,
             hf_quantizer=hf_quantizer,
-            keep_in_fp32_regex=keep_in_fp32_regex,
             device_mesh=device_mesh,
         )
 
@@ -1211,82 +1173,6 @@ def _get_dtype(
     return config, dtype, dtype_orig
 
 
-def _get_device_map(
-    model: "PreTrainedModel",
-    device_map: Optional[Union[dict, str]],
-    max_memory: Optional[dict],
-    hf_quantizer: Optional[HfQuantizer],
-    dtype: Optional[torch.dtype],
-    keep_in_fp32_regex: Optional[re.Pattern],
-) -> dict:
-    """Compute the final `device_map` to use if we passed a value in ['auto', 'balanced', 'balanced_low_0', 'sequential'].
-    Otherwise, we check for any device inconsistencies in the device_map.
-    """
-    if isinstance(device_map, str):
-        special_dtypes = {}
-        if hf_quantizer is not None:
-            special_dtypes.update(hf_quantizer.get_special_dtypes_update(model, dtype))
-        if keep_in_fp32_regex is not None:
-            special_dtypes.update(
-                {name: torch.float32 for name, _ in model.named_parameters() if keep_in_fp32_regex.search(name)}
-            )
-
-        target_dtype = dtype
-
-        if hf_quantizer is not None:
-            target_dtype = hf_quantizer.adjust_target_dtype(target_dtype)
-
-        no_split_modules = model._get_no_split_modules(device_map)
-        device_map_kwargs = {"no_split_module_classes": no_split_modules}
-
-        if "special_dtypes" in inspect.signature(infer_auto_device_map).parameters:
-            device_map_kwargs["special_dtypes"] = special_dtypes
-        elif len(special_dtypes) > 0:
-            logger.warning(
-                "This model has some weights that should be kept in higher precision, you need to upgrade "
-                "`accelerate` to properly deal with them (`pip install --upgrade accelerate`)."
-            )
-
-        if device_map != "sequential":
-            inferred_max_memory = get_balanced_memory(
-                model,
-                dtype=target_dtype,
-                low_zero=(device_map == "balanced_low_0"),
-                max_memory=max_memory,
-                **device_map_kwargs,
-            )
-        else:
-            inferred_max_memory = get_max_memory(max_memory)
-        if hf_quantizer is not None:
-            inferred_max_memory = hf_quantizer.adjust_max_memory(inferred_max_memory)
-
-        # `inferred_max_memory` contains non-reserved memory. There may be *unused* reserved memory in the GPU,
-        # which we can use to allocate parameters.
-        for device_name in inferred_max_memory:
-            if isinstance(device_name, int):  # it's a GPU device
-                if is_torch_xpu_available():
-                    unused_memory = torch.xpu.memory_reserved(device_name) - torch.xpu.memory_allocated(device_name)
-                else:
-                    unused_memory = torch.cuda.memory_reserved(device_name) - torch.cuda.memory_allocated(device_name)
-                inferred_max_memory[device_name] += unused_memory
-            # respect the `max_memory` passed by the user
-            if max_memory is not None and device_name in max_memory:
-                inferred_max_memory[device_name] = min(inferred_max_memory[device_name], max_memory[device_name])
-        device_map_kwargs["max_memory"] = inferred_max_memory
-
-        device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs)
-
-        if hf_quantizer is not None:
-            hf_quantizer.validate_environment(device_map=device_map)
-
-    elif device_map is not None:
-        tied_params = find_tied_parameters(model)
-        # check if we don't have tied param in different devices
-        check_tied_parameters_on_same_device(tied_params, device_map)
-
-    return device_map
-
-
 def _find_missing_and_unexpected_keys(
     model: "PreTrainedModel",
     original_checkpoint_keys: list[str],
@@ -4585,14 +4471,14 @@ def from_pretrained(
             # Let's make sure we don't run the init function of buffer modules
             model = cls(config, *model_args, **model_kwargs)
 
+        # Potentially upcast some modules to avoid loosing precision
+        model.upcast_modules_in_fp32(hf_quantizer, dtype)
         # Make sure to tie the weights correctly
         model.tie_weights()
 
         # make sure we use the model's config since the __init__ call might have copied it
         config = model.config
 
-        # Regex to keep a fixed dtype
-        keep_in_fp32_regex = get_keep_in_fp32_regex(model, hf_quantizer, dtype)
         if hf_quantizer is not None:  # replace module with quantized modules (does not touch weights)
             hf_quantizer.preprocess_model(
                 model=model,
@@ -4608,7 +4494,7 @@ def from_pretrained(
 
         # Prepare the full device map
         if device_map is not None:
-            device_map = _get_device_map(model, device_map, max_memory, hf_quantizer, dtype, keep_in_fp32_regex)
+            device_map = _get_device_map(model, device_map, max_memory, hf_quantizer, dtype)
 
         # restore default dtype
         if dtype_orig is not None:
@@ -4626,7 +4512,6 @@ def from_pretrained(
             disk_offload_folder=offload_folder,
             dtype=dtype,
             hf_quantizer=hf_quantizer,
-            keep_in_fp32_regex=keep_in_fp32_regex,
             device_mesh=device_mesh,
             key_mapping=key_mapping,
             weights_only=weights_only,
@@ -4806,7 +4691,6 @@ def _load_pretrained_model(
         disk_offload_folder: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
         hf_quantizer: Optional[HfQuantizer] = None,
-        keep_in_fp32_regex: Optional[re.Pattern] = None,
         device_mesh: Optional["torch.distributed.device_mesh.DeviceMesh"] = None,
         key_mapping: Optional[dict[str, str]] = None,
         weights_only: bool = True,
@@ -4873,13 +4757,6 @@ def _load_pretrained_model(
         # correctly initialize the missing (and potentially mismatched) keys
         model._initialize_missing_keys(missing_keys + mismatched_keys, is_quantized)
 
-        # Set some modules to fp32 if needed
-        if keep_in_fp32_regex is not None:
-            for name, param in model.named_parameters():
-                if keep_in_fp32_regex.search(name):
-                    # param = param.to(torch.float32) does not work here as only in the local scope.
-                    param.data = param.data.to(torch.float32)
-
         # Get reverse key mapping
         reverse_key_renaming_mapping = {v: k for k, v in key_renaming_mapping.items()}
 
@@ -4931,7 +4808,6 @@ def _load_pretrained_model(
                 reverse_key_renaming_mapping,
                 disk_offload_folder,
                 disk_offload_index,
-                keep_in_fp32_regex,
                 device_mesh,
             )
             for shard_file in checkpoint_files
@@ -4979,7 +4855,7 @@ def _load_pretrained_model(
                     if param.device.type == "meta":
                         continue
                     # Shard the param
-                    to_contiguous, casting_dtype = _infer_parameter_dtype(model, name, param, keep_in_fp32_regex)
+                    to_contiguous, casting_dtype = _infer_parameter_dtype(model, name, param)
                     shard_and_distribute_module(
                         model,
                         param.to(tp_device),
@@ -5360,6 +5236,35 @@ def train(self, mode: bool = True):
     def eval(self):
         return self.train(False)
 
+    def upcast_modules_in_fp32(self, hf_quantizer: HfQuantizer | None, dtype: torch.dtype) -> None:
+        """
+        Upcast modules defined in `_keep_in_fp32_modules` and `_keep_in_fp32_modules_strict` in fp32, if
+        `dtype` is different than fp32.
+        """
+        # If the dtype is already fp32, we can skip
+        if dtype == torch.float32:
+            return
+
+        keep_in_fp32_modules = []
+        # The _keep_in_fp32_modules flag is only used to avoid bf16 -> fp16 casting precision issues. It was introduced
+        # in case of force loading a model that should stay bf16 in fp16 (which includes a few quantizers as this is a pre-processing
+        # step for e.g. bitsandbytes). See https://github.com/huggingface/transformers/issues/20287 for details.
+        if self._keep_in_fp32_modules is not None and (
+            dtype == torch.float16 or getattr(hf_quantizer, "use_keep_in_fp32_modules", False)
+        ):
+            keep_in_fp32_modules.extend(self._keep_in_fp32_modules)
+
+        if self._keep_in_fp32_modules_strict is not None and (dtype == torch.float16 or dtype == torch.bfloat16):
+            keep_in_fp32_modules.extend(self._keep_in_fp32_modules_strict)
+
+        if len(keep_in_fp32_modules) > 0:
+            # We need to match exact layers, so we add either `.` on each side, or start/end of string
+            keep_in_fp32_regex = re.compile("|".join([rf"((^|\.){module}($|\.))" for module in keep_in_fp32_modules]))
+            for name, param in self.named_parameters():
+                if keep_in_fp32_regex.search(name):
+                    # param = param.to(torch.float32) does not work here as only in the local scope.
+                    param.data = param.data.to(torch.float32)
+
 
 PreTrainedModel.push_to_hub = copy_func(PreTrainedModel.push_to_hub)
 if PreTrainedModel.push_to_hub.__doc__ is not None:

filename: src/transformers/quantizers/base.py
patch: @@ -128,6 +128,22 @@ def adjust_target_dtype(self, dtype: "torch.dtype") -> "torch.dtype":
         """
         return dtype
 
+    def param_element_size(self, model: "PreTrainedModel", param_name: str) -> float:
+        "Return the element size (in bytes) for `param_name`."
+        if self.param_needs_quantization(model, param_name):
+            from accelerate.utils import CustomDtype
+
+            mapping = {
+                torch.int8: 1,
+                CustomDtype.INT4: 0.5,
+                CustomDtype.FP8: 1,
+                CustomDtype.INT2: 0.25,
+            }
+            # The value passed is actually not used when the method is overriden
+            if (custom_dtype := self.adjust_target_dtype(torch.float16)) in mapping:
+                return mapping[custom_dtype]
+        return model.get_parameter_or_buffer(param_name).element_size()
+
     def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
         """
         Override this method if you want to adjust the `missing_keys`.