Update PET checkpoint format

[abmazitov]

This PR fixes issue #396
Starting from now, two versions of the PET checkpoint will be saved in the end of the training
The first one will contain the data on the step of the training (i.e. optimiser and scheduler states, and also the model weights). The second one will contains only the best model state with no trainer state.
API will accept both formats and load the trainer state only if it is available.

Contributor (creator of pull-request) checklist

• Issue referenced (for PRs that solve an issue)?
• Model training works
• Model training from last checkpoint works
• Model training from best checkpoint works
• PET.restart(new_dataset_info) adapted
• Best model keeper works throughout multiple training runs
• Fine-tuning with -c works (both last and best checkpoint)
• Model loads and exports normally
• PET-MAD in checkpoints are uploaded to HF
• HF fetch works
• Documentation updated (for new features)?

---

📚 Documentation preview 📚: https://metatrain--399.org.readthedocs.build/en/399/

[frostedoyster]

It looks good to me. Just some small comments on the docs (but we can also improve them later) and potentially one bug

[frostedoyster]

What does the * stand for here? Perhaps it would be worth explaining what it is

[abmazitov]

I was referring to the fact that the checkpoint with best model is now forced to be have the best_model_ prefix. But I think here I will just leave it as best_model.ckpt, since this is the default name now.

[Luthaf]

I was referring to the fact that the checkpoint with best model is now forced to be have the best_model_ prefix

Where is this enforced? I don't see it in this PR

[abmazitov]

I was referring to the fact that the checkpoint with best model is now forced to be have the best_model_ prefix

Where is this enforced? I don't see it in this PR

It happens in the Trainer.save_checkpoint(), where I add last_checkpoint_ and best_ prefixes to the path

[frostedoyster]

Does this mean that the published checkpoint for PET-MAD will be named best_model_{something}.ckpt? Or can it be renamed and still work?

[Luthaf]

It happens in the Trainer.save_checkpoint()

Ok, but then this is only something for PET, not a general thing. This should be fine here for the checkpoint name, but I now realize that this file (which should apply to any architecture) contains a lot of PET-specific documentation (LoRA hypers, …). Should we move this elsewhere to be explicitly about PET?

[Luthaf]

Ah, there is a warning on top. Might still make sense to move the file, but this can happen later

[abmazitov]

Does this mean that the published checkpoint for PET-MAD will be named best_model_{something}.ckpt? Or can it be renamed and still work?

Nono, for sure the final name of the file can be anything, it's just the default name (i.e. the one the checkpoint will have after the training) is forced to have this prefix to distinguish between the checkpoint with only the best model, and the one with the last step of the training. You can rename the model later and use it with -c normally.

[abmazitov]

Ah, there is a warning on top. Might still make sense to move the file, but this can happen later

Maybe we can do it while working on a PR on more general and high-level fine-tuning implementation, but for now I would keep it like this