Resetting base_model breaks shared tensors for safetensors

Can someone explain how to fix a problem I am facing with safetensors saving? It seems something is being done with base_model in the original model that I need to know how to replicate. The error can be produced with the following:

model = AutoModelForAudioClassification.from_pretrained(
    "facebook/wav2vec2-base", num_labels=num_labels, label2id=label2id, id2label=id2label
)
print('Save 1')
save_file(model.state_dict(), 'temp')
print('Save 1 Complete')
model.base_model = model.base_model
print('Save 2')
save_file(model.state_dict(), 'temp')
print('Save 2 Complete')

This outputs:

Save 1                                                                                                                                                                                                                                                                         
Save 1 Complete                                                                                                                                                                                                                                                                
Save 2                                                                                                                                                                                                                                                                         
Traceback (most recent call last):        
.....                                                                                                                                                                                                                         
RuntimeError:                                                                                                                                                                                                                                                                  
            Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [{'wav2vec2.masked_spec_embed', 'base_model.masked_spec_embed'},
.....
'base_model.encoder.layers.11.final_layer_norm.bias'}].
            A potential way to correctly save your model is to use `save_model`.
            More information at https://huggingface.co/docs/safetensors/torch_shared_tensors

My only guess is that there is something being flagged with base_model to tell safetensors to ignore it, but when I reset the variable that flag is getting deleted.

1 Like

This is the easiest way to get around it, but it’s just a workaround.

–save_safetensors False

Thanks for the reply, that just causes huggingface to use torch.save instead of safetensors.save_file correct? That has worked for me in other cases but this model seems to have parameterized modules and gives the following error with torch.save

RuntimeError: Serialization of parametrized modules is only supported
through state_dict(). See:
https://pytorch.org/tutorials/beginner/saving_loading_models.html#
saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training

Do you know of any methods other than this workaround?

1 Like

that just causes huggingface to use torch.save instead of safetensors.save_file correct?

true

Perhaps this…?