How to modify model states before saving checkpoints ?

[thangld201]

I'm using Optimum and BetterTransformer in Lightning Module. When I apply BetterTransformer to the original BERT model, its state_dicts change, so if the state_dicts is saved this way there will be compatibility problems when loading it later (with from_pretrained, etc). So I want to reverse the BERT model to normal before saving the checkpoint (e.g. supposed that the BERT model is wrapped in a LightningModule).

from optimum.bettertransformer import BetterTransformer
from transformers import AutoModel

with torch.device(“cuda”):
    model = AutoModel.from_pretrained("bert-base-cased", torch_dtype=torch.float16)

model = BetterTransformer.transform(model)

# do your inference or training here

# if training and want to save the model (how to do this in Lightning ??)
model = BetterTransformer.reverse(model)
state_dict_to_save = {'state_dicts':model.state_dict()}

[tshu-w]

https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#on-save-checkpoint

[thangld201]

Thanks ! But this method is called after model.state_dict(), at which point we will already have the transformed version of BERT saved. This means we'd have to remove the transformed state_dict and then calling .reverse() to save the original BERT but this would varies per a model (backbone) basis. (a lot of code changes)

Do you know any method/callback, one that is called even before model.state_dict() so that calling .reverse() would be enough ?

[tshu-w]

Then I suggest taking a look at on_validation_end or on_validation_epoch_end, because ModelCheckpoint saves the ckpt at that time by default.