Make U matrices not persistent to reduce state_dict size.

[pfebrer]

Is your feature request related to a problem? Please describe.
We are using mace with max angular momentum 4 and correlation 3 and the checkpoint files are huge because the U matrices of the Contraction class are stored in it (~400MB when the parameters of the model only occupy ~5MB).

Describe the solution you'd like
We would like that U matrices weren't stored in the checkpoint file.

Describe alternatives you've considered
We think that passing persistent=False in this line:

mace/mace/modules/symmetric_contraction.py

Line 110 in 44b6da4

self.register_buffer(f"U_matrix_{nu}", U_matrix)

should solve the problem and doesn't cause any harm to the model. Otherwise maybe allowing the user to choose whether they are persistent or not would be nice as well.

[ilyes319]

Hey @pfebrer,

That is a reasonable solution. Could you try it and check if it solves your issues?

[pfebrer]

It indeed solved the problem, my checkpoint files size has been reduced from 680 MB to 26 MB :) I could also restart the training without problems.

But maybe in some case it is useful to store them precomputed?

[pfebrer]

I now noticed that this change breaks compatibility with loading models.

E.g.: if you load a model that stored the U matrices in a version of mace that has them set to persistent=False this will result in an error. And the same happens in the opposite case.

[peterbjorgensen]

I now noticed that this change breaks compatibility with loading models.

E.g.: if you load a model that stored the U matrices in a version of mace that has them set to persistent=False this will result in an error. And the same happens in the opposite case.

This could maybe be fixed with non-strict loading. The torch load function has a keyword strict=False, but it might be better to just break the backwards compatibility or explicitly remove the U matrices from old checkpoints.

[pfebrer]

@ilyes319 do you think you can make them persistent=False? To load an old model on the new implementation it would just be a matter of "cleaning" the checkpoint file, i.e. removing the matrices from it.

[ilyes319]

I wonder how this interacts with torchscript though and libtorch. I guess the safest would be to make an argument and keep the default to true. Would this alright?

[pfebrer]

Yes, if it can be configured from an argument of the SymmetricContraction module (not just Contraction) I think it would be fine for us 👍

[pfebrer]

Could we add this? :) (a persistent_U_matrices or something similar argument to SymmetricContraction that defaults to True)

I can submit a PR.