[Bugfix] [Offloading] Save disk-offloaded buffers, Save converted weights

src/transformers/core_model_loading.py

@@ -1508,7 +1508,7 @@ def convert_and_load_state_dict_in_model(
     return loading_info, disk_offload_index
 
 
-def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]):
+def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
     """
     Revert the conversion mapping that was used to load the model with `from_pretrained`, or the default one
     if the model was created in another way and is part of the default mappings.

src/transformers/integrations/accelerate.py

@@ -22,7 +22,7 @@
 import re
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -545,6 +545,56 @@ def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> torch
     return tensor
 
 
+def load_offloaded_checkpoint_parameters(
+    model: "PreTrainedModel",
+    param_name: str,
+    meta_state_dict: dict[str, torch.Tensor | Any] | None = None,
+) -> dict[str, torch.Tensor]:
+    """
+    Load source `param_name` from disk, if it was offloaded due to the device_map,
+    and thus lives as a meta parameter inside `model`.
+
+    This is needed when resaving a model, when some parameters were offloaded (we need to load them from disk, to
+    then resave them to disk in the correct shard...).
+
+    Example flow:
+    ```
+    source_param_name -> "experts.0.w1.weight"
+    target_param_name -> "experts.gate_up_proj"
+    loaded_state_dict -> {
+        "experts.0.w1.weight": ...
+        "experts.1.w1.weight": ...
+        ...
+        "experts.0.w3.weight": ...
+        "experts.1.w3.weight": ...
+        ...
+    }
+    ```
+
+    Args:
+        model (`PreTrainedModel`): Model containing target offloaded weight to load
+        param_name (`str`): Name of source (checkpoint) weight to load from model
+
+    Returns:
+        `dict[str, torch.Tensor]`: Loaded state dict of all weights which map to the source weight
+    """
+    from ..core_model_loading import WeightConverter, WeightRenaming, rename_source_key, revert_weight_conversion
+
+    # Convert from source key in checkpoint to target key in model
+    if meta_state_dict is None:
+        meta_state_dict = model.state_dict()  # this can be costly: please pass as arg when possible
+    renamings = [entry for entry in model._weight_conversions if isinstance(entry, WeightRenaming)]
+    converters = [entry for entry in model._weight_conversions if isinstance(entry, WeightConverter)]
+    param_name = rename_source_key(param_name, renamings, converters, model.base_model_prefix, meta_state_dict)[0]
+
+    # load parameter from model
+    tensor = load_offloaded_parameter(model, param_name)
+
+    # Convert from target key to source key(s)
+    loaded_state_dict = revert_weight_conversion(model, {param_name: tensor})
+    return loaded_state_dict
+
+
 def _init_infer_auto_device_map(
     model: nn.Module,
     max_memory: dict[int | str, int | str] | None = None,

src/transformers/modeling_utils.py

@@ -62,7 +62,7 @@
     check_and_set_device_map,
     expand_device_map,
     get_device,
-    load_offloaded_parameter,
+    load_offloaded_checkpoint_parameters,
 )
 from .integrations.deepspeed import _load_state_dict_into_zero3_model
 from .integrations.eager_paged import eager_paged_attention_forward
@@ -475,7 +475,7 @@ def remove_tied_weights_from_state_dict(
             # In offloaded cases, there may be meta tensors in the state_dict.
             # For these cases, key by the pointer of the original tensor object
             # (state_dict tensors are detached and therefore no longer shared)
-            tensor = model.get_parameter(name)
+            tensor = model.get_parameter_or_buffer(name)
             ptrs[id(tensor)].append(name)
 
         else:
@@ -3509,15 +3509,14 @@ def save_pretrained(
 
         # if any model parameters are offloaded, we need to know it for later
         is_offloaded = False
-        if (
-            hasattr(self, "hf_device_map")
-            and len(set(self.hf_device_map.values())) > 1
-            and ("cpu" in self.hf_device_map.values() or "disk" in self.hf_device_map.values())
+        if hasattr(self, "hf_device_map") and (
+            len(set(self.hf_device_map.values())) > 1 or "disk" in self.hf_device_map.values()
         ):
             is_offloaded = True
             warnings.warn(
                 "Attempting to save a model with offloaded modules. Ensure that unallocated cpu memory "
-                "exceeds the `shard_size` (50GB default)"
+                "exceeds the `shard_size` (50GB default) and/or the largest model weight size (this can "
+                "be very large for MoE models with fused experts)."
             )
 
         # Translate state_dict from smp to hf if saving with smp >= 1.10
@@ -3582,20 +3581,26 @@ def save_pretrained(
                 os.remove(full_filename)
 
         # Save the model
+        meta_state_dict = model_to_save.state_dict()
         for shard_file, tensor_names in logging.tqdm(
             state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
         ):
             filename = os.path.join(save_directory, shard_file)
             shard_state_dict = {}
-            for tensor_name in tensor_names:
+            for tensor_name in sorted(tensor_names):
                 # Get the tensor, and remove it from state_dict to avoid keeping the ref
                 tensor = state_dict.pop(tensor_name)
 
-                # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
-                # but it would otherwise not be contained in the saved shard if we were to simply move the file
-                # or something
+                # If the param was offloaded, we need to load it back onto cpu from disk to resave it.
+                # It's a strange pattern, but is necessary to ensure saving into the proper file shard
                 if is_offloaded and tensor.device.type == "meta":
-                    tensor = load_offloaded_parameter(model_to_save, tensor_name)
+                    # Note that `load_offloaded_parameter` may load multiple weights for a single tensor.
+                    # While it is possible to overload CPU memory by loading parameters in a bad order,
+                    # in practice `split_torch_state_dict_into_shards` preserves weight locality
+                    state_dict.update(
+                        load_offloaded_checkpoint_parameters(model_to_save, tensor_name, meta_state_dict)
+                    )
+                    tensor = state_dict.pop(tensor_name)
 
                 # only do contiguous after it's permuted correctly in case of TP
                 shard_state_dict[tensor_name] = tensor.contiguous()