Add latency benchmark for ML-based models

[forcadellvincent]

Context

Hi @raspstephan,

In the context of my activity, I've been given time to work on the benchmark of ML-based forecast models. WeatherBench being a great open-source project, I would like to propose my help to try to improve the benchmark.

The first idea I would like to introduce is the addition of a "Performance vs Latency" plot in the benchmark. Latency would be defined as the inference time per forecast step or for a fixed lead-time (e.g. 15 days). This is interesting for people who face restrictions on GPU resources, or those who want to perform inference at a large scale (e.g. Huge Ensembles related work), and want to know the trade-off between performance and latency.

For the figures, we could draw inspiration from the LLM community. See the example from https://huggingface.co/spaces/optimum/llm-perf-leaderboard., tab "Find Your Best Model".

Proposed steps

For one model, launch a n_steps forecast: ~10 warm-up steps, ~ 50-100 profiled steps. Latency is defined as the average time on profiled steps. Generalize to all models, and save results in the WeatherBench bucket. Then plot the results in the WeatherBench front-end.

As for the inference code, I personally find earth2studio very useful. It has a large model zoo, persistent cache management for input data, access to several datasources (WB2, CDS, ARCO, ...), and runs on Apache 2.0 License. Happy to discuss any alternatives if you have some.

Challenges foreseen

• implementation code for every model should remain in one unique package to ru(le)n them all. This way, we would reduce the implementation difference between models on the impact of the latency benchmark.
• some models currently in WeatherBench are not present in earth2studio. Might need to implement them there in the long-run.
• TPU vs GPU: how to harmonize the Latency benchmark with models originally optimized for different hardwares?
• Traditional NWP are in a category of their own, due to their CPU workloads. Likely better to discuss them in a different issue.

Hope you’ll find the contribution useful! Happy to receive feedback, hints, advice, or any clues on existing results related to this topic that I've overlooked.👍

[raspstephan]

This is a great idea but I am not sure that this fits into the scope of WB. First, as you said, for many of the models we don't have the code to run them. Second, I imagine that there is a lot of potential optimization around inference time. We, for example, run our forecasts internally on TPUs.

Maybe a more general question: AI-NWP inference is typically very fast. Do you have examples where inference speed is really a bottleneck in an operational context (maybe for creating large ensemble hindcasts I understand the problem.

[forcadellvincent]

Thank you very much for your feedback 🙏

I am totally aligned with the challenges. Inference optimization could change figures on latency by quite a margin, indeed.
As per the bottleneck of AI-NWP inference, it indeed becomes a problem when one wants to generate a certain amount of forecasts for many different initialisation dates. As per Mahesh et al. (2025), it also seems that I/O operations are probably the bottleneck for HENS, rather than GPU runtime.

About GPU runtime, Mahesh et al. (2025) give useful numbers:

Generating the [7424-members] ensemble for each initial date requires 21 TB of space and 45 min using 256
80 GB NVIDIA A100 GPUs. With 92 initial times [in June, July, August], generating HENS costs 18 432 GPU hours (256 GPUs for 3 d).

These numbers are for a certain deterministic ML model where ensembles are generated using initial condition perturbation and model perturbation. I am expecting probabilistic AI-NWP models (AIFS-Ensemble, GenCast, etc ...) to produce forecasts slower than that.

The scale of the project in Mahesh et al. (2025) is massive, but I can foresee applications in my domain that would have a rather similar setup. I feel that it is for such projects, if you add potential GPU resources restrictions and budgets, that a benchmark on latency would be more than welcomed. I understand the challenges though, and the fact that it is a bit out-of-scope of weatherbench.

On the 'GPU-operations' versus 'I/O read/write operations' bottleneck, did you also explore on your end what operations took most of the time in inference mode? On my end, I'll try to access resources and run with a few models to try to answer 👍

[raspstephan]

Thanks for your comment.

To be honest, I am not that familiar with how our models are run operationally. I suspect that this is also very Google-specific, so not sure how generalizeable these things are.