GPU parallel inference

[mchantry]

Is your feature request related to a problem? Please describe.

Models greater than the GPU memory capacity cannot be currently run in inference, whilst parallel implementations in training exist.

Describe the solution you'd like

Implement parallel inference, allowing anemoi models to be distributed across several GPUs.

Describe alternatives you've considered

No response

Additional context

No response

Organisation

ECMWF

[jswijnands]

Additional context: This also has a use case for running inference on AWS (KNMI), where it would be preferable to request multiple smaller GPU instances and combine the GPU memory through model sharding. When requesting an AWS instance with sufficient memory to run without model sharding, you typically get an instance with multiple GPUs of which only one is used

[dietervdb-meteo]

Any progress on this? We now also have a checkpoint that seems to require multiple GPU's for inference.

I see that PR #55 was closed without success?

[cathalobrien]

@dietervdb-meteo we have a PR now #108

[cathalobrien]

PR #108 has been merged, closing this issue