Multiple GPUs Inference

[BigShuiTai]

Hello, is there any way to run a inference with 2 or more GPUs?

[wesselb]

Hey @BigShuiTai! Thanks for opening an issue.

I'm guessing that you want to split a single version of the model across multiple GPUs, possibly to be able to run the model on GPU with less memory.

Unfortunately, this is not supported by the version of Aurora in this repository. Aurora, however, is just a plain PyTorch model, so model parallelism (I believe the kind that you're referring to) would be possible to implement.

[luiservela]

yes - lets make this happen!

[firatozdemir]

@wesselb is this something you are looking into in a near-future?
Was this not implemented for the original training in the paper (i.e., was Aurora simply trained with Data parallelism + activation checkpointing for the Swin3D backbone)?

[wesselb]

@firatozdemir, model parallelism is unfortunately not something we intend to implement in the near future.

For the original training, we indeed simply used DDP (FSDP, to be precise) with activation checkpointing.