From a09016d94c35c614b130db19974226781a3eff87 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 2 Jul 2026 17:07:09 +0900 Subject: [PATCH 1/5] fix save --- src/transformers/modeling_utils.py | 48 ++++++++++++++++++++++++------ 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 42cb3d974b0e..d552b4459505 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -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,21 @@ 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 many-to-one operations, it may be an issue, but usually the shards + # contain all the necessary params + 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` + weight_map.update({k: shard_file} for k in shard_state_dict.keys()) + 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 +3629,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}") From cf31d74eaa17823ba8e4b8f69a457967e5cef661 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 2 Jul 2026 17:28:54 +0900 Subject: [PATCH 2/5] fix --- src/transformers/modeling_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index d552b4459505..5818634f3fa0 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3608,13 +3608,14 @@ def save_pretrained( 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 many-to-one operations, it may be an issue, but usually the shards - # contain all the necessary params + # 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` - weight_map.update({k: shard_file} for k in shard_state_dict.keys()) + weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys()) except Exception: raise RuntimeError( "We could not revert some weight conversions because of offlading, and several weights needed for a single " From 78377b46a22eca51d976803b6a16d79e219de95c Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 2 Jul 2026 18:35:29 +0900 Subject: [PATCH 3/5] fix --- src/transformers/modeling_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 5818634f3fa0..a7d61329c6b8 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3615,7 +3615,8 @@ def save_pretrained( 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` - weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys()) + if state_dict_split.is_sharded: + weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys()) except Exception: raise RuntimeError( "We could not revert some weight conversions because of offlading, and several weights needed for a single " From 4b0bc7aea940a03ebb79f1e9537d66cb3eda5cc5 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 2 Jul 2026 18:44:41 +0900 Subject: [PATCH 4/5] ignore stupid ty --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a7d61329c6b8..980b63f747c9 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3616,7 +3616,7 @@ def save_pretrained( 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()) + 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 " From 1ee8e252d8eff47fbd0f47ce465861bcd73476fc Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 2 Jul 2026 18:50:57 +0900 Subject: [PATCH 5/5] add buffers --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 980b63f747c9..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: