Repository: optuna/optuna
Pull request: Add device argument to BoTorchSampler
<!-- Thank you for creating a pull request! In general, we merge your pull request after it gets two or more approvals. To proceed to the review process by the maintainers, please make sure that the PR meets the following conditions: (1) it passes all CI checks, and (2) it is neither in the draft nor WIP state. If you wish to discuss the PR in the draft state or need any other help, please mention the Optuna development team in the PR. -->

## Motivation
<!-- Describe your motivation why you will submit this PR. This is useful for reviewers to understand the context of PR. -->

Resolve https://github.com/optuna/optuna/issues/4096.

## Description of the changes
<!-- Describe the changes in this PR. -->

Add `device` argument to `BoTorchSampler` and send input data of BoTorch's candidate functions to the device.


Naive benchmark results on Google colab were 1 min (GPU) vs 2 min (CPU). The code I used as follows:


```python
import torch
import optuna
from optuna.integration import BoTorchSampler


def o(t):
    for i in range(30):
        t.suggest_int(f"x_{i}", -100, 100)
        t.suggest_categorical(f"y_{i}", list(range(-100, 100)))

    return t.suggest_float("z", 1e-4, 10000, log=True)

# GPU
s = optuna.create_study(sampler=BoTorchSampler(device=torch.device("cuda:0")))
s.optimize(o, n_trials=50) # about 1 min

# CPU
s = optuna.create_study(sampler=BoTorchSampler(device=torch.device("cpu")))
s.optimize(o, n_trials=50) # about 2 min
```

## [updated]

The colab notebook is available at https://colab.research.google.com/drive/1PMfHH2ZlZ_FL1aFez5V56XzUmAFM2eil?usp=sharing.

Problem statement:
Support GPU in BoTorchSampler
### Motivation

Currently, BoTorchSampler only supports CPU computation, while the BoTorch backend supports cuda accelaration.

### Description

Add `device: optional[torch.device] = None` argument to `BoTorchSampler.__init__` that can be used in the following way.
```python
sampler = BoTorchSampler()  # CPU
sampler = BoTorchSampler(device=torch.device("cuda:0"))  # GPU
```

### Alternatives (optional)

_No response_

### Additional context (optional)

We don't require a test of the GPU case for now because it is difficult to run in the current CI.

Hints from issue discussion:

Commit message: Add test of device argument

Patch to review:
diff --git a/optuna/integration/botorch.py b/optuna/integration/botorch.py
index 36084b26655..c5ec11b0f40 100644
--- a/optuna/integration/botorch.py
+++ b/optuna/integration/botorch.py
@@ -328,8 +328,6 @@ def _get_default_candidates_func(
         return qei_candidates_func
 
 
-# TODO(hvy): Allow utilizing GPUs via some parameter, not having to rewrite the callback
-# functions.
 @experimental_class("2.4.0")
 class BoTorchSampler(BaseSampler):
     """A sampler that uses BoTorch, a Bayesian optimization library built on top of PyTorch.
@@ -386,6 +384,9 @@ class BoTorchSampler(BaseSampler):
             conditional.
         seed:
             Seed for random number generator.
+        device:
+            A ``torch.device`` to store input and output data of BoTorch. Please set a CUDA device
+            if you fasten sampling.
     """
 
     def __init__(
@@ -406,6 +407,7 @@ def __init__(
         n_startup_trials: int = 10,
         independent_sampler: Optional[BaseSampler] = None,
         seed: Optional[int] = None,
+        device: Optional["torch.device"] = None,
     ):
         _imports.check()
 
@@ -417,6 +419,7 @@ def __init__(
 
         self._study_id: Optional[int] = None
         self._search_space = IntersectionSearchSpace()
+        self._device = device if device is not None else torch.device("cpu")
 
     def infer_relative_search_space(
         self,
@@ -505,11 +508,11 @@ def sample_relative(
                     "constraints. Constraints passed to `candidates_func` will contain NaN."
                 )
 
-        values = torch.from_numpy(values)
-        params = torch.from_numpy(params)
+        values = torch.from_numpy(values).to(self._device)
+        params = torch.from_numpy(params).to(self._device)
         if con is not None:
-            con = torch.from_numpy(con)
-        bounds = torch.from_numpy(bounds)
+            con = torch.from_numpy(con).to(self._device)
+        bounds = torch.from_numpy(bounds).to(self._device)
 
         if con is not None:
             if con.dim() == 1:
@@ -546,7 +549,7 @@ def sample_relative(
                 f"{candidates.size(0)}, bounds: {bounds.size(1)}."
             )
 
-        return trans.untransform(candidates.numpy())
+        return trans.untransform(candidates.cpu().numpy())
 
     def sample_independent(
         self,
diff --git a/tests/integration_tests/test_botorch.py b/tests/integration_tests/test_botorch.py
index efdd97b6e12..4ff1ec88fd0 100644
--- a/tests/integration_tests/test_botorch.py
+++ b/tests/integration_tests/test_botorch.py
@@ -410,3 +410,22 @@ def test_call_after_trial_of_independent_sampler() -> None:
     ) as mock_object:
         study.optimize(lambda _: 1.0, n_trials=1)
         assert mock_object.call_count == 1
+
+
+@pytest.mark.parametrize("device", [None, torch.device("cpu"), torch.device("cuda:0")])
+def test_device_argument(device: Optional[torch.device]) -> None:
+
+    sampler = BoTorchSampler(device=device)
+    if not torch.cuda.is_available() and sampler._device.type == "cuda":
+        pytest.skip(reason="GPU is unavailable.")
+
+    def objective(trial: Trial) -> float:
+        return trial.suggest_float("x", 0.0, 1.0)
+
+    def constraints_func(trial: FrozenTrial) -> Sequence[float]:
+        x0 = trial.params["x"]
+        return [x0 - 0.5]
+
+    sampler = BoTorchSampler(constraints_func=constraints_func, n_startup_trials=1)
+    study = optuna.create_study(sampler=sampler)
+    study.optimize(objective, n_trials=3)


Code review comments:
diff_hunk: @@ -417,6 +419,7 @@ def __init__(
 
         self._study_id: Optional[int] = None
         self._search_space = IntersectionSearchSpace()
+        self._device = device if device is not None else torch.device("cpu")
path: optuna/integration/botorch.py
text: @user1:
It can be simpler.
```suggestion
        self._device = device or torch.device("cpu")
```