CUDA Out of Memory – Aurora ERA5 Deployment

[gingerlolipop]

I'm deploying Aurora using the ERA5 example: https://microsoft.github.io/aurora/example_era5.html.​

The deployment fails with the following error: RuntimeError: CUDA out of memory. Tried to allocate XX GiB. GPU 0 has a total capacity of XX GiB of which XX MiB is free...

I've attempted to reduce the rollout steps, but the issue persists. I understand that Azure VMs offer higher RAM options, but they are cost-prohibitive for my current needs.​

So, what is the minimum GPU VRAM required to successfully run this Aurora example? Are there recommended configurations or optimizations to reduce GPU memory usage?​

Any guidance would be appreciated.

[wesselb]

Hey @gingerlolipop! Thanks for opening an issue.

To run the full model, you'll need approximately 40 GB of GPU memory. You might be able to bring that down a little with optimisations.

[gingerlolipop]

@wesselb I see. Thanks! The best GPU I can get is V100 +32GB GPU memory. But I can obtain 4 at a time. Does this model support multi GPU? Otherwise I will crop the map to make it smaller!

[wesselb]

The implementation unfortunately currently does not support model parallelism, so you cannot split the model over multiple GPUs. This is something that you'd have to implement yourself.

[BigShuiTai]

Try to use PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:32. I can run the Aurora finetuned model using MI50 32GB GPU with this setting.

[wesselb]

@gingerlolipop, let me know if @BigShuiTai's suggestion works! If it does, perhaps we can add your suggestion to the docs, @BigShuiTai. :)

[leonardledy]

@gingerlolipop Since you have at least 3 GPUs, you can try manually allocating each block of the model (encoder-processor-decoder) to a different GPU.
It's simple but it worked fine for me on 0.25° data with 3 GPUs of 24GB. I don't think that the number of rollout steps matters much here since for every prediction you only use the past two predictions/initial conditions.

Hope it helps!