Wrapping architectures to ensure consistent behavior and reduce workload on devs and maintainers.

[pfebrer]

This comes after discussions on how to best implement #1209 . The idea there is that we can not apply hooks on architectures externally because at inference time we need to apply them before the scaler and additive models, whose application was always inside the forward of the models.

Even without hooks, the fact that the architecture's model handled its own preprocessing lead to very strange patterns, e.g.:

• The model is passed a dataset_info on init, but then it is asked to produce the data in a different format (specially bad for atomic basis stuff).
• Convoluted workflows to do finetuning, since the model has to be aware of how it will be trained next, instead of just receiving simple instructions like "remove this output" or "add this output".

In this PR, the trainer has:

• A Trainer.setup() method that is in charge of setting up the model (already knowing how it will be trained). It can then set up all the preprocessing, pass the model the true layouts that it will be asked to produce, and then wrap it in a MetatrainWrapper that will handle applying the reverse processing.
• A Trainer.restart() method that is in charge of restarting the preprocessing pipelines, coming up with the new dataset_info to pass to the model and then calling model.restart().

For now I have implemented this idea in PET, and even for PET still needs some polishing. Before proceeding I would like to get some feedback from @PicoCentauri and @Luthaf to see if I have to go in some other direction.

So as a summary of the changes:

• There is a metatrain/utils/wrapper.py that implements the MetatrainWrapper.
• Adapted PET Trainer to follow the described approach, moved some things from the Model to the Trainer.
• The scaler and composition needed some slight modifications on their restart functions.
• Modified the metatrain/cli/train.py script to use this new flow.

---

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

[Luthaf]

Looks good overall, seems that it will simplify things quite a bit.

I'm not sure about having MetatrainWrapper follow the same API as all other architecture. It does not have to, and we can maybe make our job even more easier by not forcing it to.

IMO we do not need dataset_info in __init__, checkpoint handling/loading could likely be simplified, and maybe other things could be removed/simplified

[Luthaf]

not sure if this is the best way, these don't really feel like hypers. I'm fine if we make MetatrainWrapper "special", and not force it to follow the same API as an architecture

[pfebrer]

Yes they are way different from the rest of hypers in the sense that serializing them and having them in a file would be crazy. So the main idea of following the same interface exactly is to make this indistinguishable for metatomic. Which part of the API does metatomic use/need?

[Luthaf]

Mainly the forward function signature. The functions for NL requests and extra inputs are recursively explored, and will be found on any submodule.

The __init__ and class attributes can be whatever you want!

The full documentation is https://docs.metatensor.org/metatomic/latest/torch/reference/models/export.html#metatomic.torch.ModelInterface

[pfebrer]

Ok, do you think this class can inherit from metatomic's ModelInterface?

[Luthaf]

it can if you want, but it does not have to. ModelInterface is only for documentation

[Luthaf]

naming suggestion:

Suggested change

	class MetatrainWrapper(ModelInterface[WrapperHypers]):
	class MetatrainModel(ModelInterface[WrapperHypers]):

[pfebrer]

Hmm so you are thinking that all architectures should be forced to use the wrapper?

[Luthaf]

Would be a lot easier, no? And worst case the wrapper does nothing if an architecture does not want a scaler/additive models

[pfebrer]

Yeah maybe it indeed simplifies things, let me try and see what happens