Best practices for distributing NN model weights

[fcasson]

If I understand correctly, the model weights for the new QLKNN are in one of these files ?

https://github.com/google-deepmind/fusion_surrogates/tree/main/fusion_surrogates/qlknn/models

As a community, we have growing collection of surrogate models and the distribution and packaging of their weights could do with some standardisation / best practices. Maybe you have thought about this more than the wider fusion community...

Some questions

• Are model weights best distributed via git, git lfs, or something else ?
• I don't want to be continually downloading unrelated git histories and working trees that I don't really need
• Is it good to use git submodules to pull these weights into my code that uses them? Upside is good provenance tracking, downside is continually cloning other stuff we don't need, and submodules are not well suited to managing optional dependencies
• Do we all agree on ONNX format as best for moving between frameworks?

[fcasson]

After thinking about this for a while, we decided

1. We really have to seperate NN weights distribution from super project cloning, so it won't be in our source tree (git submodules are out for this), but the old QLKNN is a bit different because we also have to build some of the source there.
2. This means it doesn't matter so much how they are distributed. git is okay, and git lfs is probably better.
3. However they are distributed it is really helpful to have clear versioning of models, and this requires are a bit of thought if too many different models are in a single repo.
4. To support wider usage, it's helpful if onnx versions can be provided.

[jcitrin]

@hamelphi thoughts here?

Regarding the onnx version, we will reprioritize that such that you can get QLKNN_7_11 into JINTRAC/HFPS

[hamelphi]

Apologies for the delay in response

Yes, model weights are stored in fusion_surrogates/qlknn/models/qlknn_7_11.qlknn. The format is serialized and can be loaded using qlknn_model.py

We have given some thought about standardization of surrogate models. And this library is an attempt to provide helpers and distribute surrogate models weights and inference code in a somewhat standardized way. However, different frameworks have different requirements, so it is difficult to find something that works for everyone (e.g. Python vs. Fortran, Linux vs. MacOs vs. Windows).

For instance, we found that providing the path to model weights is often cumbersome to do in a portable way, as environment variables can cause silent bugs or unexpected behaviour, while flags can be annoying to specify for every run. We attempted to solve this by having a model "registry". Having the inference code independent of the weights can also be problematic (e.g. QLKNN-hyper) as implementation of the inference code need to be validated against external code.

Though the solutions we came up with are suited to our specific needs, we may be a bit blind to what problems are present in external workflows. I am open to suggestions on how we can improve accessibility to the models.

1-2. If using python, an alternative to git is to use pip install which doesn't involve git history or test data, only what is selected to be packaged, meaning the weights and inference code. It is possible to provide options for the installation so we can have fine-grain control on which model/version we want to install. I don't have experience with git lfs personally. This may be worth it for large models, but for small models like QLKNN_7_11 (< 1Mb), pure git storage or pip packaging is fine.

3: Agreed that versioning is important. If models are included as part of a library, like QLKNN_7_11, Models should be versioned within the library, this is what is attempted in the registry, the idea being that we wouldn't overwrite the old version as we update the model, but rather add a new version. We would then provide pip install options to select which model version(s) to install to avoid installing unneeded deprecated models.

4: Yes, we are actively looking into solutions to provide an ONNX version of the model to allow wider use of the model outside Python workflows. Unfortunately, JAX is not officially supported by ONNX (https://onnx.ai/supported-tools.html#buildModel), but there are open-source options that we are looking into.

[fcasson]

1: We all seem to be using git, with git lfs useful for large models.

4: Thanks, looks like the community is standardising on ONNX for now, so we postpone our attempts to use StableHLO:

• https://github.com/ukaea/tglfnn-ukaea
• https://github.com/ProjectTorreyPines/TurbulentTransport.jl/blob/master/utilities/README_onnx_to_pytorch_fortran.md
• Add qlknn_7_11 onnx model export #25

[fcasson]

To summarise what we've learnt so far from TGLFNN-UKAEA about weights distribution:

• git LFS is okay, but not on github due to very restrictive rate limits on the free tier. Hugging face also has git interface, but it seems rate limits much more generous since it is targeted at ML -> We may move to that.
• pypi looks to be okay (for python based projects) if total download stays under 10GB.
• CI pipelines downloading the weights many times still require a bit of thought to minimize unnecessary downloads - caching, avoid parallelising tests over the same weights etc?

And about framework conversions:

• Best maintained onnx to jax (and reverse) conversion process still seems to be an open question for the longer term ?

@theo-brown may also want to weigh in