Repository: huggingface/transformers
Pull request: Offloading need to add the prefix into the offload_index

# What does this PR do?



Changed files:
filename: src/transformers/integrations/accelerate.py
patch: @@ -466,10 +466,10 @@ def expand_device_map(device_map, param_names):
 
 
 def accelerate_disk_offload(
+    model: "PreTrainedModel",
     disk_offload_folder: str | None,
     checkpoint_files: list[str] | None,
     device_map: dict,
-    expected_keys: list[str],
     sharded_metadata: dict | None,
     dtype: torch.dtype | None,
     weight_mapping=None,
@@ -493,7 +493,8 @@ def accelerate_disk_offload(
     # In this case, the offload index is simply the existing safetensors (except if using custom weight loading
     # Operation, e.g. the MoE models, where we need to resave the weights that were changed at loading time)
     if is_offloaded_safetensors:
-        param_device_map = expand_device_map(device_map, expected_keys)
+        meta_state_dict = model.state_dict()
+        param_device_map = expand_device_map(device_map, meta_state_dict.keys())
         str_dtype = str(dtype).replace("torch.", "") if dtype is not None else "float32"
         if sharded_metadata is None:
             weight_map = dict.fromkeys(safe_open(checkpoint_files[0], framework="pt").keys(), checkpoint_files[0])
@@ -502,7 +503,9 @@ def accelerate_disk_offload(
             weight_map = {k: os.path.join(folder, v) for k, v in sharded_metadata["weight_map"].items()}
 
         # Update the weight names according to the `weight_mapping`
-        weight_renaming_map = {rename_source_key(k, renamings, [])[0]: k for k in weight_map}
+        weight_renaming_map = {
+            rename_source_key(k, renamings, [], model.base_model_prefix, meta_state_dict)[0]: k for k in weight_map
+        }
 
         # Prepare the index using existing safetensors files
         disk_offload_index = {

filename: src/transformers/integrations/peft.py
patch: @@ -299,13 +299,14 @@ def load_adapter(
             renamings = [entry for entry in key_mapping 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:

filename: src/transformers/modeling_utils.py
patch: @@ -4089,10 +4089,10 @@ def _load_pretrained_model(
         # Prepare parameters offloading if needed
         if device_map is not None and "disk" in device_map.values():
             disk_offload_index = accelerate_disk_offload(
+                model,
                 disk_offload_folder,
                 checkpoint_files,
                 device_map,
-                expected_keys,
                 sharded_metadata,
                 dtype,
                 weight_mapping,