fix

Author: Cyrilvallez

src/transformers/core_model_loading.py

@@ -30,6 +30,7 @@
 
 import torch
 
+from .integrations.accelerate import offload_weight
 from .integrations.tensor_parallel import ALL_PARALLEL_STYLES, DTensor, Replicate, TensorParallelLayer
 from .utils import is_torch_greater_or_equal, logging
 
@@ -344,7 +345,7 @@ def dot_natural_key(s: str):
 
 @contextmanager
 def log_to_misc(
-    layer_name: str,
+    full_param_name: str,
     misc: MutableMapping[str, str],
     extras: Any = None,
     op: Union[list[ConversionOps], ConversionOps, None] = None,
@@ -368,30 +369,30 @@ def _format_op_name(curr_op: Union[list[ConversionOps], ConversionOps, None]) ->
         if isinstance(extras, tuple) and len(extras) == 2:
             values, target_keys = extras
             descriptor = f"{op_name} " if op_name else ""
-            misc[layer_name] = (
+            misc[full_param_name] = (
                 f"{e}\nError: {descriptor}on tensors destined for {target_keys}. Ckpt contains: {len(values[0])}"
             )
         elif isinstance(extras, str):
             suffix = f" via {op_name}" if op_name else ""
-            misc[layer_name] = f"{e}\nError{suffix} when processing parameter {extras}"
+            misc[full_param_name] = f"{e}\nError{suffix} when processing parameter {extras}"
         elif extras is None and op_name:
-            misc[layer_name] = f"{op_name}: {e}"
+            misc[full_param_name] = f"{op_name}: {e}"
         else:
-            misc[layer_name] = f"{extras} |Error: {e}"
+            misc[full_param_name] = f"{extras} |Error: {e}"
         raise SkipLayer()
 
 
 def set_param_for_module(
     model: PreTrainedModel,
-    layer_name: str,
+    full_param_name: str,
     param_value: torch.Tensor,
     mismatch_keys: MutableSet[tuple[str, torch.Size, torch.Size]],
     missing_keys: MutableSet[str],
     misc: MutableMapping[str, Any],
     distributed_operation: Optional[TensorParallelLayer],
 ):
-    with log_to_misc(layer_name, misc, layer_name):
-        module_path, _, param_name = layer_name.rpartition(".")
+    with log_to_misc(full_param_name, misc, full_param_name):
+        module_path, _, param_name = full_param_name.rpartition(".")
         module_obj = model.get_submodule(module_path) if module_path else model
         param_value = param_value[0] if isinstance(param_value, list) else param_value[...]
         ref = getattr(module_obj, param_name)
@@ -414,9 +415,9 @@ def set_param_for_module(
                 param_value = torch.nn.Parameter(param_value, requires_grad=param_value.is_floating_point())
 
         # Remove from missing keys (it's either mismatched, or all good)
-        missing_keys.discard(layer_name)
+        missing_keys.discard(full_param_name)
         if ref is not None and ref.shape != param_value.shape:
-            mismatch_keys.add((layer_name, param_value.shape, ref.shape))
+            mismatch_keys.add((full_param_name, param_value.shape, ref.shape))
             module_obj.param_name._is_hf_initialized = False  # Needs to be initialized
         else:
             param_value._is_hf_initialized = True  # super important otherwise _init_weight re-initi if bias is missing
@@ -439,6 +440,8 @@ def convert_and_load_state_dict_in_model(
     device_map: dict | None = None,
     dtype_plan: dict | None = None,
     device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
+    disk_offload_index: dict | None = None,
+    disk_offload_folder: str | None = None,
 ):
     """
     Convert a state dict according to a weight mapping (one WeightConverter per glob pattern),
@@ -448,6 +451,7 @@ def convert_and_load_state_dict_in_model(
     prefix = model.base_model_prefix
     tp_plan = tp_plan or {}  # {glob_pattern: plan_obj_or_key}
     device_map = device_map or {}  # {exact_target_key: device}
+    device_map_regex = "|".join([re.escape(k) for k in sorted(device_map.keys(), reverse=True)])
     dtype_plan = dtype_plan or {}  # {glob_pattern: dtype}
     weight_mapping = weight_mapping or {}  # {glob_pattern: WeightConverter}
     meta_model_state_dict = model.state_dict()
@@ -533,7 +537,7 @@ def convert_and_load_state_dict_in_model(
                     shard_index,
                 )
 
-        if future is None:  # If not TP, async materialize the tensors. TODO handle disk offload?
+        if future is None:
             future = spawn_materialize(thread_pool, tensor, _dtype)
         entry.collected_tensors[target_key].setdefault(converter_key, []).append(future)
 
@@ -546,29 +550,29 @@ def convert_and_load_state_dict_in_model(
             group = by_conversion_pattern.pop(key)
             converter = group.weight_converter
             operations = converter.operations if isinstance(converter.operations, list) else [converter.operations]
-            for layer_name, tensors_for_this_layer in group.collected_tensors.items():
+            for full_param_name, tensors_for_this_layer in group.collected_tensors.items():
                 pbar.update(1)
-                pbar.set_postfix({"Materializing param": layer_name})
+                pbar.set_postfix({"Materializing param": full_param_name})
                 pbar.refresh()
-                concrete_target_keys = layer_name.split("|")
+                concrete_target_keys = full_param_name.split("|")
                 try:
                     if bool(set(concrete_target_keys) - unexpected_keys):
-                        with log_to_misc(layer_name, misc):
+                        with log_to_misc(full_param_name, misc):
                             values = [[k.result() for k in inner] for inner in tensors_for_this_layer.values()]
 
                         for op in operations:
-                            with log_to_misc(layer_name, misc, (values, concrete_target_keys), operations):
+                            with log_to_misc(full_param_name, misc, (values, concrete_target_keys), operations):
                                 values = op.convert(values, model.config)
 
                         values = [values] if not isinstance(values, list) else values
-                        with log_to_misc(layer_name, misc, (values, concrete_target_keys), operations):
+                        with log_to_misc(full_param_name, misc, (values, concrete_target_keys), operations):
                             realized_value = {
                                 k: t for k, t in zip(concrete_target_keys, values) if k not in unexpected_keys
                             }
 
                         for k in list(realized_value.keys()).copy():
                             if op := converter.quantization_operation:
-                                with log_to_misc(layer_name, misc, op=op):
+                                with log_to_misc(full_param_name, misc, op=op):
                                     realized_value.update(
                                         op.convert(
                                             {k: realized_value.pop(k)}, quant_config=quantizer.quantization_config
@@ -578,15 +582,26 @@ def convert_and_load_state_dict_in_model(
                         for k, output_value in realized_value.items():
                             for src in converter.source_keys:  # what should happen to k when we meet k at saving
                                 inverse_converters[k] = {src: converter}
-                            set_param_for_module(
-                                model,
-                                k,
-                                output_value,
-                                mismatch_keys,
-                                missing_keys,
-                                misc,
-                                converter.distributed_operation,
-                            )
+
+                            param_device = device_map[re.search(device_map_regex, k).group()]
+                            # Offloading support
+                            if param_device == "disk":
+                                missing_keys.discard(k)
+                                # If not already offloaded, or if we applied any special Operation, we need to re-save
+                                if k not in disk_offload_index or len(operations) > 0:
+                                    disk_offload_index = offload_weight(
+                                        output_value, k, disk_offload_folder, disk_offload_index
+                                    )
+                            else:
+                                set_param_for_module(
+                                    model,
+                                    k,
+                                    output_value,
+                                    mismatch_keys,
+                                    missing_keys,
+                                    misc,
+                                    converter.distributed_operation,
+                                )
 
                 except SkipLayer:
                     continue

src/transformers/integrations/accelerate.py

@@ -508,7 +508,7 @@ def accelerate_disk_offload(
         os.makedirs(disk_offload_folder, exist_ok=True)
     is_offloaded_safetensors = checkpoint_files is not None and checkpoint_files[0].endswith(".safetensors")
 
-    # In this cause, the offload index is simply the existing safetensors (except if using custom weight loading
+    # 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)

src/transformers/modeling_utils.py

@@ -59,7 +59,6 @@
     check_and_set_device_map,
     expand_device_map,
     init_empty_weights,
-    offload_weight,
 )
 from .integrations.deepspeed import _load_state_dict_into_zero3_model
 from .integrations.eager_paged import eager_paged_attention_forward
@@ -4268,6 +4267,8 @@ def _load_pretrained_model(
                 device_map,
                 model.dtype_plan,
                 device_mesh,
+                disk_offload_index,
+                disk_offload_folder,
             )
 
             # finally close all opened file pointers
@@ -4319,13 +4320,6 @@ def _load_pretrained_model(
                         device_mesh,
                     )
 
-        # If the model parameters were changed during loading (i.e. any custom Ops on the weights), we need to resave them
-        # for offloading
-        if device_map is not None and "disk" in device_map.values():
-            for name, param in model.state_dict().items():
-                if name not in disk_offload_index:
-                    disk_offload_index = offload_weight(param, name, disk_offload_folder, disk_offload_index)
-
         log_state_dict_report(
             model=model,
             pretrained_model_name_or_path=pretrained_model_name_or_path,