Inference of NequIP and Allegro implemented into CP2K and question about metadata

[gabriele16]

Hi,

Thanks for the development of such impressive work and for making it open-source.

We have implemented inference of NequIP as well as Allegro into the CP2K molecular mechanics package, FIST. This was recently merged into the CP2K development branch and will soon appear in CP2K v2023.2.

If possible, I would like to ask to add the options model_dtype and default_dtype to the metadata, for a future release of NequIP and/or Allegro where running with mixed precision is enabled (as shown in the preprint https://doi.org/10.48550/arXiv.2304.10061). In this way, their values can be read directly when when parsing the deployed .PTH file, instead of looking at the .YAML file itself.

For example, in pair_nequip/pair_nequip.cpp or pair_allegro/pair_allegro.cpp this would allow to read model_dtype and default_dtype as extra metadata key/value pairs in the following way:

std::unordered_map<std::string, std::string> metadata = {
    {"config", ""},
    {"nequip_version", ""},
    {"r_max", ""},
    {"n_species", ""},
    {"type_names", ""},
    {"_jit_bailout_depth", ""},
    {"_jit_fusion_strategy", ""},
    {"allow_tf32", ""},
    {"model_dtype",""},
    {``default_dtype'', ""}
  };

  model = torch::jit::load(std::string(arg[2]), device, metadata);
  model.eval();

We would then adapt this to the CP2K/LibTorch interface, instead of parsing the config value to extract default_dtype as currently done, which is not ideal, see #cp2k/cp2k#2786 (comment)

Thanks again.

[Linux-cpp-lisp]

Hi @gabriele16 ,

Wow!! Exciting to see this and thank you for your efforts in this direction and reaching out.

On develop, aka 0.6.*---which will be released to main soon and which should be your only development target for new code---I think this is already implemented: https://github.com/mir-group/nequip/blob/develop/nequip/scripts/deploy.py#L314

I also have logic in there for getting a (usually, not always) correct default from the metadata for old models that don't store it: https://github.com/mir-group/nequip/blob/develop/nequip/scripts/deploy.py#L127

Would you be willing to have a quick Zoom call at some point? I'd love to chat and learn more about this effort. My email is in my GitHub profile, please feel free to reach out.

[gabriele16]

Hi @Linux-cpp-lisp

Thanks a lot for the quick reply and for the nice feedback :) Ah, I should have checked the develop branch better...It's great to know that it's already implemented and that there is a way to make sure that old models work correctly nevertheless.

Of course, I will be very happy to chat over Zoom! Exchanging info on this and other features/topics would be good also to make sure that we keep the CP2K implementation correctly up-to-date with the latest nequip/allegro developments. I will write soon.

[Linux-cpp-lisp]

Hi @gabriele16 ,

Great! Looking forward to hear from you by email and discuss, thanks!