diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 6c9fe775c8b5..105f934a287d 100644 --- a/src/transformers/core_model_loading.py +++ b/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. diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index 63a3ce22ea55..0c258482b5fd 100644 --- a/src/transformers/integrations/accelerate.py +++ b/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, diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 42cb3d974b0e..02b0fcb58304 100644 --- a/src/transformers/modeling_utils.py +++ b/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()