Trainer's `save_model` isn't saving the entire state_dict and is only saving the embedding/encoder
I currently have a very basic model that consist of a pre-trained backbone model with a MLP head. Think something like:class Model(PreTrainedModel):
    def __init__(self, config):
        self.backbone = AutoModel.from_pretrained(config.model_name_or_path)
        self.mlp = MLPLayer()

    ...I’m currently using the Trainer object to save my model. Specifically, the code that I’m using has aself.save_model(output_dir)inside the Training loop. This saves the checkpoint as a safetensors object.When I try to load it usingmodel = AutoModel.from_pretrained(PATH_TO_SAFETENSORS)I’m noticing that the keys for the MLP layer are not there and only the keys for the backbone model’s embedding and encoder layers are there.I took a look at the source code forsave_model, which seems to be using the_savemethod, and don’t see any reason why the MLP layers shouldn’t be saved. Both the_saveandsave_pretrainedmethods use thestate_dictwhich contains the MLP layer’s weights.Is there anything that I may be missing or may have configured incorrectly? Thanks.
2024-01-02T04:22:51Z
Well the mistake was embarrassingly simple. It turns out that if I don’t initialize the exact model object, then only the keys that exist will be loaded.That means that since I was doingmodel = AutoModel.from_pretrained(PATH_TO_SAFETENSORS)and notmodel = Model.from_pretrained(PATH_TO_SAFETENSORS)the code was (correctly) only loading thestate_dictfor the embedding and encoder layers, since my own MLP layer doesn’t exist inAutoModel.I do think that a warning message of some type would be a bit helpful, but just leaving this here for other to find.