Feature/mapper chunking

[japols]

Describe your changes

This PR adds a chunking strategy for the GraphTransformerMapperBlock to reduce peak memory usage during inference.

During inference, the max memory usage is dominated by large peaks in the the encoder and decoder. Chunking the edges and computing the GraphTransformer forward pass sequentially nicely flattens these peaks without significant increases in runtime. The following projection MLPs are also computed sequentially in (node) chunks.
Here are some results for a single-gpu, 240h-leadtime inference run on n320, hidden grid o48:

num_chunks	Peak Memory usage (GB)	Time (s)
1	33.45	159.13
2	29.81	159.45
4	24.85	160.77
8	22.91	160.84
16	22.08	164.38

The number of chunks during inference is handled by an environment variable ANEMOI_NUM_CHUNKS_INFERENCE which seems to be the only practical solution currently as we are currently not able to pass options to the model during inference. It falls back to no chunking i.e. num_chunks=1 if not set.
I am happy about any suggestions/alternative solutions.

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update*

*I am not sure how/where to document the usage of the environment variable, the goal would be to find a better long-time solution at some point

Checklist before requesting a review

• I have performed a self-review of my code
• My code follows the style guidelines of this project
• I have commented my code, particularly in hard-to-understand areas
• I have updated the documentation and docstrings to reflect the changes
• I have added tests that prove my fix is effective or that my feature works
• I have ensured that the code is still pip-installable after the changes and runs
• I have not introduced new dependencies in the inference partion of the model*
• I have ran this on single GPU
• I have ran this on multi-GPU or multi-node
• I have ran this to work on LUMI (or made sure the changes work independently.)
• I have ran the Benchmark Profiler against the old version of the code

*technically the environment variable is a new dependency in inference

Tag possible reviewers

@ssmmnn11 @JesperDramsch @theissenhelen @gmertes

---

📚 Documentation preview 📚: https://anemoi-models--46.org.readthedocs.build/en/46/

[FussyDuck]

All committers have signed the CLA.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.84%. Comparing base (57ea969) to head (b6e34b9).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop      #46   +/-   ##
========================================
  Coverage    99.84%   99.84%           
========================================
  Files           23       23           
  Lines         1277     1301   +24     
========================================
+ Hits          1275     1299   +24     
  Misses           2        2           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[gmertes]

Regarding the environment variable:

We will come up with an interface to pass options to the model during inference. In the meantime, I suggested to Jan we work with an environment variable because this does not require a code change across the anemois.

When we have this options interface, we can load inference settings from a predefined place (the user config, command line, etc). While the environment variable is not a long-term solution, we can keep it as "one of the sources to look for options", lowest in the priority list.

We will delegate the options handling to anemoi-inference, which then passes the option values to the model. So down the line this env var will probably move to anemoi-inference, into a general inference options handler. But I suggest we deal with the later, so we don't block Jan's work.

[ssmmnn11]

Great work!