Repository: huggingface/transformers
Pull request: fix torch_dtype, contiguous, and load_state_dict regression

# What does this PR do ? 

This PR fixes the regression we have from https://github.com/huggingface/transformers/pull/36335:
- the model should be in fp32 if no torch_dtype is passed 
- param should be contiguous if the old_param is
- change back `load_state_dict` map_location default to "cpu" + couples of small fixes -> fixes CI ! 
- switch from `get_submodule_get_module_from_name` to `get_module_name` is it also handles the case where the param name reference to a direct param from the model. 

Should fixes tests in peft and in the quantization CI

Perf with TP 2 GPUS:

```
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:02<00:00,  1.34it/s]
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:03<00:00,  1.33it/s]
Model loading time: 5.91 seconds
Model loading time: 5.77 seconds
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                 model_load        63.70%        3.572s       100.00%        5.607s        5.607s       0.000us         0.00%     773.210ms     773.210ms           8 b     -14.96 Gb       7.97 Gb     -29.92 Gb             1
                                aten::copy_        14.63%     820.125ms        45.88%        2.572s       3.593ms     773.210ms       100.00%     949.505ms       1.326ms           0 b      -2.25 Gb           0 b           0 b           716
                                   aten::to         0.03%       1.862ms        29.27%        1.641s       1.124ms       0.000us         0.00%     773.210ms     529.596us         272 b           0 b       7.97 Gb           0 b          1460
                             aten::_to_copy         0.08%       4.381ms        29.23%        1.639s       2.788ms       0.000us         0.00%     773.210ms       1.315ms         272 b           0 b       7.97 Gb           0 b           588
                            cudaMemcpyAsync        14.14%     792.725ms        14.14%     792.725ms       2.715ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           292
                           aten::contiguous         0.01%     321.399us        13.54%     759.069ms      11.860ms       0.000us         0.00%       0.000us       0.000us       2.25 Gb           0 b           0 b           0 b            64
                                aten::clone         0.02%       1.045ms        13.53%     758.747ms      11.855ms       0.000us         0.00%       0.000us       0.000us       2.25 Gb           0 b           0 b           0 b            64
                             cudaMemGetInfo         4.80%     269.310ms         4.87%     273.148ms     273.148ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1
                               aten::detach         0.14%       7.882ms         1.13%      63.286ms      19.062us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          3320
                                     detach         0.86%      48.257ms         0.99%      55.334ms      19.504us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          2837
                           _FromTorchTensor         0.38%      21.113ms         0.44%      24.477ms     108.785us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           225
                                aten::empty         0.16%       8.994ms         0.24%      13.663ms       8.971us       0.000us         0.00%       0.000us       0.000us      47.12 Gb      47.12 Gb      29.92 Gb      29.92 Gb          1523
                      cudaStreamSynchronize         0.23%      12.651ms         0.23%      12.651ms      43.324us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           292
                        aten::empty_strided         0.17%       9.539ms         0.18%       9.957ms      16.297us       0.000us         0.00%       0.000us       0.000us         272 b         272 b       7.97 Gb       7.97 Gb           611
                                 aten::set_         0.14%       7.715ms         0.14%       7.715ms       6.594us       0.000us         0.00%       0.000us       0.000us     -29.92 Gb     -29.92 Gb           0 b           0 b          1170
                                 aten::view         0.10%       5.355ms         0.10%       5.355ms       3.852us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          1390
                                 cudaMalloc         0.09%       4.963ms         0.09%       4.963ms       2.481ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             2
                               aten::select         0.07%       3.705ms         0.08%       4.329ms       7.426us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           583
                                   Resource         0.07%       3.731ms         0.07%       3.731ms     932.860us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4
                                 aten::item         0.02%       1.049ms         0.06%       3.617ms       6.214us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           582
                              aten::view_as         0.02%       1.331ms         0.06%       3.364ms      14.951us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           225
                                aten::slice         0.05%       3.006ms         0.06%       3.346ms       6.484us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           516
                  aten::_local_scalar_dense         0.05%       2.568ms         0.05%       2.568ms       4.412us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           582
                              aten::reshape         0.02%       1.028ms         0.04%       2.218ms       3.811us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           582
                           aten::as_strided         0.02%       1.010ms         0.02%       1.010ms       0.868us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          1163
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 5.607s
Self CUDA time total: 773.210ms

Loading took 5.909772157669067 seconds
Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                 model_load        66.96%        3.596s       100.00%        5.371s        5.371s       0.000us         0.00%     769.832ms     769.832ms           8 b     -14.96 Gb       7.97 Gb     -29.92 Gb             1
                                aten::copy_        15.43%     828.993ms        48.13%        2.585s       3.611ms     769.832ms       100.00%     944.418ms       1.319ms           0 b      -2.25 Gb           0 b           0 b           716
                                   aten::to         0.04%       1.920ms        30.65%        1.646s       1.128ms       0.000us         0.00%     769.832ms     527.282us         272 b           0 b       7.97 Gb           0 b          1460
                             aten::_to_copy         0.08%       4.544ms        30.61%        1.644s       2.796ms       0.000us         0.00%     769.832ms       1.309ms         272 b           0 b       7.97 Gb           0 b           588
                            cudaMemcpyAsync        14.69%     788.871ms        14.69%     788.871ms       2.702ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           292
                           aten::contiguous         0.01%     310.266us        14.30%     768.117ms      12.002ms       0.000us         0.00%       0.000us       0.000us       2.25 Gb           0 b           0 b           0 b            64
                                aten::clone         0.01%     695.911us        14.29%     767.807ms      11.997ms       0.000us         0.00%       0.000us       0.000us       2.25 Gb           0 b           0 b           0 b            64
                               aten::detach         0.15%       8.003ms         1.14%      61.366ms      18.484us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          3320
                                     detach         0.86%      46.253ms         0.99%      53.289ms      18.784us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          2837
                           _FromTorchTensor         0.39%      21.086ms         0.46%      24.558ms     109.145us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           225
                                aten::empty         0.18%       9.418ms         0.26%      13.991ms       9.186us       0.000us         0.00%       0.000us       0.000us      47.12 Gb      47.12 Gb      29.92 Gb      29.92 Gb          1523
                      cudaStreamSynchronize         0.25%      13.234ms         0.25%      13.234ms      45.322us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           292
                        aten::empty_strided         0.18%       9.468ms         0.18%       9.875ms      16.162us       0.000us         0.00%       0.000us       0.000us         272 b         272 b       7.97 Gb       7.97 Gb           611
                             cudaMemGetInfo         0.16%       8.779ms         0.16%       8.779ms       8.779ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1
                                 aten::set_         0.16%       8.561ms         0.16%       8.561ms       7.317us       0.000us         0.00%       0.000us       0.000us     -29.92 Gb     -29.92 Gb           0 b           0 b          1170
                                 aten::view         0.10%       5.547ms         0.10%       5.547ms       3.990us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          1390
                                 cudaMalloc         0.09%       4.880ms         0.09%       4.880ms       2.440ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             2
                               aten::select         0.07%       3.773ms         0.08%       4.339ms       7.442us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           583
                              aten::view_as         0.02%       1.330ms         0.06%       3.472ms      15.430us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           225
                                aten::slice         0.06%       2.984ms         0.06%       3.338ms       6.470us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           516
                                 aten::item         0.02%       1.074ms         0.05%       2.607ms       4.479us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           582
                              aten::reshape         0.02%       1.050ms         0.04%       2.227ms       3.826us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           582
                  aten::_local_scalar_dense         0.03%       1.532ms         0.03%       1.532ms       2.633us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b           582
                           aten::as_strided         0.02%     964.845us         0.02%     964.845us       0.830us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b          1163
                           aten::empty_like         0.01%     290.848us         0.01%     688.124us      10.752us       0.000us         0.00%       0.000us       0.000us       2.25 Gb           0 b           0 b           0 b            64
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 5.371s
Self CUDA time total: 769.832ms
```



Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -67,6 +67,7 @@
     translate_to_torch_parallel_style,
 )
 from .quantizers import AutoHfQuantizer, HfQuantizer
+from .quantizers.quantizers_utils import get_module_from_name
 from .safetensors_conversion import auto_conversion
 from .utils import (
     ACCELERATE_MIN_VERSION,
@@ -536,11 +537,11 @@ def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True):
 def load_state_dict(
     checkpoint_file: Union[str, os.PathLike],
     is_quantized: bool = False,
-    map_location: Optional[Union[str, torch.device]] = "meta",
+    map_location: Optional[Union[str, torch.device]] = "cpu",
     weights_only: bool = True,
 ):
     """
-    Reads a `safetensor` or a `.bin` checkpoint file into `meta` if requested.
+    Reads a `safetensor` or a `.bin` checkpoint file. We load the checkpoint on "cpu" by default.
     """
     if checkpoint_file.endswith(".safetensors") and is_safetensors_available():
         with safe_open(checkpoint_file, framework="pt") as f:
@@ -771,6 +772,7 @@ def _load_state_dict_into_meta_model(
     unexpected_keys=None,  # passing `unexpected` for cleanup from quantization items
     device_mesh=None,
     shard_file=None,
+    weights_only=True,
 ):
     """
     This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its
@@ -800,7 +802,15 @@ def _load_state_dict_into_meta_model(
     if shard_file.endswith(".safetensors"):
         file_pointer = safe_open(shard_file, framework="pt", device=tensor_device)
     else:
-        bin_state_dict = load_state_dict(shard_file, map_location="cpu")
+        map_location = "cpu"
+        if (
+            device_map is not None
+            and hf_quantizer is not None
+            and hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO
+            and hf_quantizer.quantization_config.quant_type in ["int4_weight_only", "autoquant"]
+        ):
+            map_location = torch.device([d for d in device_map.values() if d not in ["cpu", "disk"]][0])
+        bin_state_dict = load_state_dict(shard_file, map_location=map_location, weights_only=weights_only)
 
     error_msgs = []
 
@@ -822,23 +832,36 @@ def _load_state_dict_into_meta_model(
             if shard_file.endswith(".safetensors")
             else bin_state_dict[serialized_param_name]
         )
+
+        # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model, and which
+        # uses `param.copy_(input_param)` that preserves the contiguity of the parameter in the model.
+        # Reference: https://github.com/pytorch/pytorch/blob/db79ceb110f6646523019a59bbd7b838f43d4a86/torch/nn/modules/module.py#L2040C29-L2040C29
+
+        old_param = model
+        splits = fixed_param_name.split(".")
+        for split in splits:
+            # We shouldn't hit the default value unless for quant methods like hqq that modifies expected_keys.
+            old_param = getattr(old_param, split, None)
+            if old_param is None:
+                break
+
+        if not isinstance(old_param, (torch.nn.Parameter, torch.Tensor)):
+            old_param = None
+
         # We convert floating dtypes to the `dtype` passed except for float8_e4m3fn type. We also want to keep the buffers/params
         # in int/uint/bool and not cast them.
         param_casting_dtype = None
         is_param_float8_e4m3fn = is_torch_e4m3fn_available and empty_param.dtype == torch.float8_e4m3fn
-
-        if dtype is not None and empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn:
-            if (
-                keep_in_fp32_modules is not None
-                and keep_in_fp32_modules.search(fixed_param_name)
-                and dtype == torch.float16
-            ):
+        if empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn:
+            if keep_in_fp32_modules is not None and keep_in_fp32_modules.search(fixed_param_name):
                 param_casting_dtype = torch.float32
-            else:
+            elif dtype is not None:
                 param_casting_dtype = dtype
+            elif old_param is not None:
+                param_casting_dtype = old_param.dtype
 
         if device_mesh is not None:  # In this case, the param is already on the correct device!
-            module_to_tp, param_type = find_submodule_and_param_name(model, fixed_param_name)
+            module_to_tp, param_type = get_module_from_name(model, fixed_param_name)
             current_module_plan = None
             full_tp_plan_ = "|".join(full_tp_plan.keys()).replace("*", "[0-9]+")
             if plan := re.search(full_tp_plan_, fixed_param_name):
@@ -859,8 +882,10 @@ def _load_state_dict_into_meta_model(
                 else:
                     param = param[rank * (row // device_mesh.size()) : (rank + 1) * (row // device_mesh.size()), :]
                     shard = Shard(0)
-                if param_casting_dtype is not None and param_casting_dtype != empty_param.dtype:
+                if param_casting_dtype is not None:
                     param = param.to(param_casting_dtype)
+                if old_param.is_contiguous():
+                    param = param.contiguous()
                 local_parameter = DTensor.from_local(
                     param,
                     device_mesh=device_mesh,
@@ -873,9 +898,18 @@ def _load_state_dict_into_meta_model(
                 output_fn = partial(tp_layer._prepare_output_fn, tp_layer.output_layouts, tp_layer.use_local_output)
                 distribute_module(module_to_tp, device_mesh, None, input_fn, output_fn)
             else:
-                module_to_tp.load_state_dict({param_type: param[:]}, strict=False, assign=True)
+                param = param[:]
+                if old_param is not None and old_param.is_contiguous():
+                    param = param.contiguous()
+                module_to_tp.load_state_dict({param_type: param}, strict=False, assign=True)
 
         else:
+            param = param[:]
+            if param_casting_dtype is not None:
+                param = param.to(param_casting_dtype)
+            if old_param is not None and old_param.is_contiguous():
+                param = param.contiguous()
+
             if device_map is None:
                 param_device = "cpu"
             else:
@@ -887,9 +921,9 @@ def _load_state_dict_into_meta_model(
 
             if param_device == "disk":
                 if not is_safetensors:
-                    offload_index = offload_weight(param[:], fixed_param_name, offload_folder, offload_index)
+                    offload_index = offload_weight(param, fixed_param_name, offload_folder, offload_index)
             elif param_device == "cpu" and state_dict_index is not None:
-                state_dict_index = offload_weight(param[:], fixed_param_name, state_dict_folder, state_dict_index)
+                state_dict_index = offload_weight(param, fixed_param_name, state_dict_folder, state_dict_index)
             elif (
                 not is_quantized
                 or (not hf_quantizer.requires_parameters_quantization)
@@ -906,23 +940,21 @@ def _load_state_dict_into_meta_model(
             ):
                 if is_fsdp_enabled():
                     param_device = "cpu" if is_local_dist_rank_0() else "meta"
-                module, param_type = find_submodule_and_param_name(model, fixed_param_name)
-                if param_casting_dtype is not None and param_casting_dtype != empty_param.dtype:
-                    param = param[:].to(param_casting_dtype)
+                module, param_type = get_module_from_name(model, fixed_param_name)
                 module.load_state_dict(
-                    {param_type: param[:].to(param_device)},
+                    {param_type: param.to(param_device)},
                     strict=False,
                     assign=True,
                 )
             else:
                 hf_quantizer.create_quantized_param(
-                    model, param[:], fixed_param_name, param_device, state_dict, unexpected_keys
+                    model, param, fixed_param_name, param_device, state_dict, unexpected_keys
                 )
                 # For quantized modules with FSDP/DeepSpeed Stage 3, we need to quantize the parameter on the GPU
                 # and then cast it to CPU to avoid excessive memory usage on each GPU
                 # in comparison to the sharded model across GPUs.
                 if is_fsdp_enabled() or is_deepspeed_zero3_enabled():
-                    module, param_type = find_submodule_and_param_name(model, fixed_param_name)
+                    module, param_type = get_module_from_name(model, fixed_param_name)
                     value = getattr(module, param_type)
                     param_to = "cpu"
                     if is_fsdp_enabled() and not is_local_dist_rank_0():
@@ -4203,7 +4235,9 @@ def from_pretrained(
                             elif not is_sharded:
                                 torch_dtype = get_state_dict_dtype(state_dict)
                             else:
-                                one_state_dict = load_state_dict(resolved_archive_file[0], weights_only=weights_only)
+                                one_state_dict = load_state_dict(
+                                    resolved_archive_file[0], map_location="meta", weights_only=weights_only
+                                )
                                 torch_dtype = get_state_dict_dtype(one_state_dict)
                                 del one_state_dict  # free CPU memory
                             logger.info(
@@ -4848,7 +4882,7 @@ def _load_pretrained_model(
         else:
             folder = None
 
-        model.expected_keys = expected_keys
+        model_to_load.expected_keys = expected_keys
         if device_map is not None:
             expanded_device_map = expand_device_map(device_map, original_loaded_keys, start_prefix)
             if hf_quantizer is None:
@@ -4907,6 +4941,7 @@ def _load_pretrained_model(
                     unexpected_keys=unexpected_keys,
                     device_mesh=device_mesh,
                     resolved_archive_file=resolved_archive_file,
+                    weights_only=weights_only,
                 )
             else:
                 # We need to read the state dict as it is meta otherwise
@@ -4957,16 +4992,8 @@ def _load_pretrained_model(
                 # Skip the load for shards that only contain disk-offloaded weights when using safetensors for the offload.
                 if shard_file in disk_only_shard_files:
                     continue
-                map_location = None
-                if (
-                    device_map is not None
-                    and hf_quantizer is not None
-                    and hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO
-                    and hf_quantizer.quantization_config.quant_type in ["int4_weight_only", "autoquant"]
-                ):
-                    map_location = torch.device([d for d in device_map.values() if d not in ["cpu", "disk"]][0])
                 state_dict = load_state_dict(
-                    shard_file, is_quantized=is_quantized, map_location=map_location, weights_only=weights_only
+                    shard_file, is_quantized=is_quantized, map_location="meta", weights_only=weights_only
                 )
 
                 # Mistmatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not
@@ -5006,6 +5033,7 @@ def _load_pretrained_model(
                             unexpected_keys=unexpected_keys,
                             device_mesh=device_mesh,
                             shard_file=shard_file,
+                            weights_only=weights_only,
                         )
                         error_msgs += new_error_msgs
                 else: