emdgroup / baybe

Bayesian Optimization and Design of Experiments
https://emdgroup.github.io/baybe/
Apache License 2.0
265 stars 42 forks source link

GPU support #356

Open mhrmsn opened 2 months ago

mhrmsn commented 2 months ago

As I recently looked into this and after discussion with @AdrianSosic and @Scienfitz, here's my observations:

I think there are two convenient ways to support GPUs, either by allowing the user to use torch.set_default_device("cuda") or by adding a configuration variable within the package, something like baybe.options.device = "cuda". I only tested the first one, I think the latter is a bit more complex, but would potentially allow to have different parts of code in the overall BayBe workflow to use different devices (this may be useful if you are generating embeddings from a torch-based neural network for use in BayBe etc.)

When experimenting with torch.set_default_device("cuda"), I noticed that the devices for the tensors are not consistently set in BayBE. For either solution I think these points would need to be addressed:

mhrmsn commented 2 months ago

Last point is now addressed in this PR: pytorch/botorch#2502

AdrianSosic commented 2 months ago

Hi @mhrmsn, thanks for investigating and raising the botorch issue 🥇

Solid GPU support is definitely something we should target and in particular also test – which brings up yet another point to think about. That involves both figuring out what are reasonable GPU tests and also how/where we can run them (AFAIK Github now offers GPU runners, but I have no clue if we can get easy access).

Regarding the other points you mentions: