qKnowledgeGradient CPU & GPU usage

[ruhanaazam]

Hi, I am interested in using qKnowledgeGradient in a simple BO loop. When I run the code below, the CPU usage on my machine is extremely high while GPU is fairly low. Is there a way to make running qKnowledgeGradient more computationally efficient?

import torch
from tqdm import tqdm
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_mll
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.models.transforms.outcome import Standardize, Normalize
from botorch.optim import optimize_acqf
from botorch.acquisition import qKnowledgeGradient,


device = torch.device(f"cuda:0" if torch.cuda.is_available() else "CPU")

train_X = torch.rand(10, 2, dtype=torch.float64, device=device) * (2149)
train_Y = lookup(train_X)

# main bayes_opt training loop
for i in tqdm(range(100)):
    gp = SingleTaskGP(
        train_X,
        train_Y,
        outcome_transform=Standardize(1),
        input_transform=Normalize(train_X.shape[-1]),
    )

    mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
    fit_gpytorch_mll(mll)

    acq = qKnowledgeGradient(gp, num_fantasies=num_fantasies)
    new_X, acq_value = optimize_acqf(
        acq,
        q=1,
        bounds=bounds,
        num_restarts=num_restarts,
        raw_samples=raw_samples,
    )
    new_Y = lookup(new_X)

    # add new candidate
    train_X = torch.cat([train_X, new_X])
    train_Y = torch.cat([train_Y, new_Y])

System information

Botorch: 0.8.3
GPyTorch: 1.10
Pytorch: 2.0.1
Ubuntu 20.04
Python: 3.10

[saitcakmak]

Hi @ruhanaazam. The GPU will generally speed things up for larger tensor operations. With BO, the sizes of the underlying tensors are typically rather small, due to the small data size. This limits how much the GPU is utilized, as well as any performance benefits from using the GPU.

[Balandat]

The main use case for GPU in BoTorch today are Hypervolume-based acquisition functions, where the computation of the (expected) hypervolume involves rather large tensors and can benefit a lot from using the GPU. See https://proceedings.neurips.cc/paper/2020/hash/6fec24eac8f18ed793f5eaad3dd7977c-Abstract.html for some context on this.

[saitcakmak]

I'll close this issue and keep it around as a discussion for future discoverability.

[sangttruong]

qKG uses 3000% CPU while qEI only uses 100%, and our GPU servers don't have much CPU. Would you recommend running KG (and generally lookahead acquisition function) on the CPU server? I tried this, and it is slower than a GPU-based job about 15 times.

[RylanSchaeffer]

@saitcakmak I don't think this problem is solved. It seems like running this job consumes a lot of CPU and very little GPU, but if one runs without a GPU, the runtime is 15x longer.

That seems like it doesn't add up?

[saitcakmak]

Can you profile where that runtime increase is coming from? The GPU is used for evaluating the acquisition function & the model predictions within, so it may speed things up there. The optimization part of optimize_acqf will use CPU only, since that is what the scipy optimizers we use support. It could be that CPU shows high utilization due to the optimizer usage while GPU speeds up the acquisition evaluation despite having low utilization (due to relatively small tensors).