diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 42cb3d974b0e..8a89fc0e6841 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -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: @@ -3538,8 +3538,11 @@ def save_pretrained( # Remove tied weights as safetensors do not handle them state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save) - # Revert all renaming and/or weight operations - if save_original_format and not _hf_peft_config_loaded: + # Revert all renaming and/or weight operations. In general, due to potential many-weights-to-one conversion patterns, + # we need to revert the whole state_dict at once to make sure all weights are available. For offloaded models though, + # some weights are on meta device, so we need to first load them back into cpu then convert (and we do it later inside a + # given shard, to avoid blowing cpu memory since offloading means constrained resources) + if save_original_format and not is_offloaded and not _hf_peft_config_loaded: state_dict = revert_weight_conversion(model_to_save, state_dict) # Shard the model if it is too big. @@ -3553,13 +3556,6 @@ def save_pretrained( state_dict_split = split_torch_state_dict_into_shards( state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size ) - # Save index if sharded - index = None - if state_dict_split.is_sharded: - index = { - "metadata": {"total_parameters": self.num_parameters(), **state_dict_split.metadata}, - "weight_map": state_dict_split.tensor_to_filename, - } # Clean the folder from a previous save for filename in os.listdir(save_directory): @@ -3581,6 +3577,17 @@ def save_pretrained( ): os.remove(full_filename) + # The weight_map may change compared to state_dict_split.tensor_to_filename due to revert weight conversions + weight_map = None + if state_dict_split.is_sharded: + # For offloaded weights, we will convert each shard later, so the weight names will change and we will fill + # the weight_map as we get them + weight_map = ( + state_dict_split.tensor_to_filename + if not (is_offloaded and save_original_format and not _hf_peft_config_loaded) + else {} + ) + # Save the model for shard_file, tensor_names in logging.tqdm( state_dict_split.filename_to_tensors.items(), desc="Writing model shards" @@ -3600,6 +3607,23 @@ def save_pretrained( # only do contiguous after it's permuted correctly in case of TP shard_state_dict[tensor_name] = tensor.contiguous() + # As explained above, for offloaded scenarios, weight format could not be reverted before due to meta weights, + # so do it now after they were loaded onto cpu. For one-weight-to-many operations, it may be an issue, but usually the shards + # contain all the necessary params, except if we are quite unlucky on the sharding. The failure surface is (very few models + # with one-weight-to-many + offloading to disk + unlucky sharding), so it will almost never happen + if is_offloaded and save_original_format and not _hf_peft_config_loaded: + try: + shard_state_dict = revert_weight_conversion(model_to_save, shard_state_dict) + # Save the weight_map, since some names etc may have changed due to conversion compared to initial `state_dict_split` + if state_dict_split.is_sharded: + weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys()) # ty: ignore[unresolved-attribute] + except Exception: + raise RuntimeError( + "We could not revert some weight conversions because of offlading, and several weights needed for a single " + "conversion operation living in different shard files. Try reducing `max_shard_size` a bit, or worst case " + "set `save_original_format=False`." + ) + # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL, # so it's not possible for now.... # Write the shard to disk @@ -3607,6 +3631,14 @@ def save_pretrained( # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM) del shard_state_dict + # Save index if sharded + index = None + if state_dict_split.is_sharded: + index = { + "metadata": {"total_parameters": self.num_parameters(), **state_dict_split.metadata}, + "weight_map": weight_map, + } + if index is None: path_to_weights = os.path.join(save_directory, weights_name) logger.info(f"Model weights saved in {path_to_weights}")