diff --git a/docs/examples/single_objective_bayes_opt/fast_model_eval.ipynb b/docs/examples/single_objective_bayes_opt/fast_model_eval.ipynb index ef9a5a260..d9535cb1e 100644 --- a/docs/examples/single_objective_bayes_opt/fast_model_eval.ipynb +++ b/docs/examples/single_objective_bayes_opt/fast_model_eval.ipynb @@ -108,10 +108,10 @@ "metadata": {}, "outputs": [], "source": [ - "# first call does tracing\n", + "# first call compiles the model\n", "fig, ax = X.generator.visualize_model(\n", " n_grid=30,\n", - " model_compile_mode=\"trace\",\n", + " model_compile_mode=\"inductor\",\n", " output_names=[\"y2\"],\n", " show_acquisition=False,\n", " show_samples=False,\n", @@ -125,10 +125,10 @@ "metadata": {}, "outputs": [], "source": [ - "# second invocation uses a pre-traced model and is a tiny bit faster (most of walltime is used by plotting)\n", + "# second invocation can reuse PyTorch's compilation cache and is a tiny bit faster (most wall time is used by plotting)\n", "fig, ax = X.generator.visualize_model(\n", " n_grid=30,\n", - " model_compile_mode=\"trace\",\n", + " model_compile_mode=\"inductor\",\n", " output_names=[\"y2\"],\n", " show_acquisition=False,\n", " show_samples=False,\n", diff --git a/xopt/generators/bayesian/utils.py b/xopt/generators/bayesian/utils.py index 38ce6520c..797553e39 100644 --- a/xopt/generators/bayesian/utils.py +++ b/xopt/generators/bayesian/utils.py @@ -332,94 +332,6 @@ def validate_turbo_controller_center(generator: Generator) -> None: ) -class MeanVarModelWrapper(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.model = model - - def forward(self, x): - output_dist = self.model(x) - return output_dist.mean, output_dist.variance - - -class MeanVarModelWrapperPosterior(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.model = model - - def forward(self, x): - output_dist = self.model.posterior(x) - return output_dist.mean, output_dist.variance - - -def torch_trace_gp_model( - model: Model, - vocs: VOCS, - tkwargs: dict, - posterior: bool = True, - grad: bool = False, - batch_size: int = 1, - verify: bool = False, -) -> torch.jit.ScriptModule: - """ - Trace a GPyTorch model using torch.jit.trace. Note that resulting object will return mean and variance directly, - NOT a multivariate normal. - - Parameters - ---------- - model : Model - The GPyTorch model to compile. - vocs : VOCS - VOCS - tkwargs : dict - The keyword arguments for the torch tensor. - posterior : bool, optional - If True, prime the model by using posterior method, otherwise call directly (this invokes gpytorch posterior). - grad : bool, optional - If True, use gradient context, otherwise use no gradient context. - batch_size : int, optional - The batch size for the input tensor for tracing, by default 1. - verify : bool, optional - If True, request that torch verify the trace by comparing to eager mode, by default False. - """ - if isinstance(model, ModelListGP): - raise ValueError( - "ModelListGP is not supported for JIT tracing - use individual models" - ) - rand_point = random_inputs(vocs)[0] - rand_vec = torch.stack( - [rand_point[k] * torch.ones(batch_size) for k in vocs.variable_names], dim=1 - ) - test_x = rand_vec.to(**tkwargs) - # test_x_1 = test_x[:1,...] - - gradctx = nullcontext() if grad else torch.no_grad() - model.eval() - with gradctx, gpytorch.settings.fast_pred_var(), gpytorch.settings.trace_mode(): - if posterior: - pred = model.posterior(test_x) - traced_model = torch.jit.trace( - MeanVarModelWrapperPosterior(model), test_x, check_trace=False - ) - traced_model = torch.jit.optimize_for_inference(traced_model) - else: - pred = model(test_x) - traced_model = torch.jit.trace( - MeanVarModelWrapper(model), test_x, check_trace=False - ) - traced_model = torch.jit.optimize_for_inference(traced_model) - if verify: - traced_mean, traced_var = traced_model(test_x) - assert torch.allclose(pred.mean, traced_mean, rtol=0), ( - f"JIT traced mean != original {pred.mean=} {traced_mean=}" - ) - assert torch.allclose(pred.variance, traced_var, rtol=0), ( - f"JIT traced variance != original: {pred.variance=} {traced_var=}" - ) - - return traced_model.to(**tkwargs) - - def torch_compile_gp_model( model: Model, vocs: VOCS, @@ -457,64 +369,30 @@ def torch_compile_gp_model( ) test_x = rand_vec.to(**tkwargs) - gradctx = nullcontext if grad else torch.no_grad() - # TODO: check if gpytorch trace mode faster + gradctx = nullcontext() if grad else torch.no_grad() with gradctx, gpytorch.settings.fast_pred_var(): model.eval() if posterior: pred = model.posterior(test_x) - traced_model = torch.compile( + compiled_model = torch.compile( model, backend=backend, mode=mode, dynamic=None ) - mvn = traced_model.posterior(test_x) + mvn = compiled_model.posterior(test_x) else: pred = model(test_x) - traced_model = torch.compile( + compiled_model = torch.compile( model, backend=backend, mode=mode, dynamic=None ) - mvn = traced_model(test_x) - traced_mean, traced_var = mvn.mean, mvn.variance - assert torch.allclose(pred.mean, traced_mean, rtol=0), ( - f"Compiled mean != original {pred.mean=} {traced_mean=}" + mvn = compiled_model(test_x) + compiled_mean, compiled_var = mvn.mean, mvn.variance + assert torch.allclose(pred.mean, compiled_mean, rtol=0), ( + f"Compiled mean != original {pred.mean=} {compiled_mean=}" ) - assert torch.allclose(pred.variance, traced_var, rtol=0), ( - f"Compiled variance != original: {pred.variance=} {traced_var=}" + assert torch.allclose(pred.variance, compiled_var, rtol=0), ( + f"Compiled variance != original: {pred.variance=} {compiled_var=}" ) - return traced_model - - -def torch_trace_acqf( - acq: AcquisitionFunction, vocs: VOCS, tkwargs: dict -) -> torch.jit.ScriptModule: - """ - Trace an acquisition function using torch.jit.trace. - - Parameters - ---------- - acq : AcquisitionFunction - The acquisition function to trace. - vocs : VOCS - VOCS - tkwargs : dict - The keyword arguments for the torch tensor. - """ - # Note that this is very fragile for when we mix q=1 and q>1 because tensors ndims changes - rand_point = random_inputs(vocs)[0] - rand_vec = torch.stack( - [rand_point[k] * torch.ones(1) for k in vocs.variable_names], dim=1 - ) - test_x = rand_vec.to(**tkwargs) - test_x = test_x.unsqueeze(-2) - with gpytorch.settings.fast_pred_var(), gpytorch.settings.trace_mode(): - # Need dummy evaluation to set caches - acq(test_x.clone().detach()) - saqcf = torch.jit.trace( - acq, - example_inputs=test_x.clone().detach(), - check_trace=False, - ) - return saqcf + return compiled_model def torch_compile_acqf( @@ -543,10 +421,9 @@ def torch_compile_acqf( verify : bool, optional If True, do the verification vs eager mode. """ - # TODO: check if trace mode better # NOTE: is verify is False, you need to ensure tensors are copied before calling # or RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run - with gpytorch.settings.fast_pred_var(), gpytorch.settings.trace_mode(): + with gpytorch.settings.fast_pred_var(): # assume that only a few shapes will happen - batch=1 and batch=nsamples saqcf = torch.compile(acq, backend=backend, mode=mode, dynamic=False) if verify: diff --git a/xopt/generators/bayesian/visualize.py b/xopt/generators/bayesian/visualize.py index 5b29773dd..84d639605 100644 --- a/xopt/generators/bayesian/visualize.py +++ b/xopt/generators/bayesian/visualize.py @@ -36,7 +36,7 @@ ) from .objectives import feasibility -from .utils import torch_compile_gp_model, torch_trace_gp_model +from .utils import torch_compile_gp_model # Little helper class, which is only used as a type. DType = TypeVar("DType") @@ -195,7 +195,9 @@ def visualize_model( exponentiate : bool, optional Flag to exponentiate acquisition function before plotting. model_compile_mode : str, optional - Compilation mode for the model. If None (default), the model is not compiled. + Compilation mode for the model. Use ``"inductor"`` to compile the model + with PyTorch's default compiler backend. If None (default), the model is + not compiled. tkwargs: dict, optional kwargs for torch tensor creation interactive: bool, optional @@ -268,6 +270,7 @@ def visualize_model( n_grid=n_grid, idx=idx, interactive=interactive, + model_compile_mode=model_compile_mode, ) ax[i, 0].set_xlabel(None) if show_acquisition: @@ -437,6 +440,7 @@ def plot_model_prediction( color: str = "C0", axis: Optional[Axes] = None, interactive: bool = False, + model_compile_mode: Optional[str] = None, ) -> Axes: """Displays the GP model prediction for the selected output. @@ -475,6 +479,8 @@ def plot_model_prediction( Whether to enable picker functionality for samples in the subplots. emit_warning : bool, optional Whether to emit a Python warning when acquisition evaluation is skipped for contextual axes. + model_compile_mode : str, optional + See eponymous parameter of :func:`visualize_model`. Returns ------- @@ -503,6 +509,7 @@ def plot_model_prediction( model=model, vocs=vocs, include_prior_mean=show_prior_mean or requires_prior_mean, + model_compile_mode=model_compile_mode, ) if len(variable_names) == 1: var_name = variable_names[0] @@ -1418,7 +1425,9 @@ def _get_model_predictions( include_prior_mean : bool, optional Whether to include the prior mean in the predictions. model_compile_mode: str, optional - Compilation mode for the model. If None (default), the model is not compiled. + Compilation mode for the model. Use ``"inductor"`` to compile the model + with PyTorch's default compiler backend. If None (default), the model is + not compiled. _ Returns @@ -1429,27 +1438,11 @@ def _get_model_predictions( gp = model.models[vocs.output_names.index(output_name)] # input_mesh = input_mesh.unsqueeze(-2) with torch.no_grad(), gpytorch.settings.fast_pred_var(): - if model_compile_mode == "trace": - if hasattr(model, "_jit"): - jitgp = model._jit - else: - jitgp = torch_trace_gp_model( - gp, - vocs, - {"device": input_mesh.device}, - posterior=True, - grad=False, - batch_size=input_mesh.shape[-1], - ) - model._jit = jitgp - mean, std = jitgp(input_mesh) - posterior_mean = mean.detach().squeeze().cpu().numpy() - posterior_std = torch.sqrt(std.detach()).squeeze().cpu().numpy() - elif model_compile_mode == "inductor": - jitgp = torch_compile_gp_model( + if model_compile_mode == "inductor": + compiled_gp = torch_compile_gp_model( gp, vocs, {"device": input_mesh.device}, posterior=True, grad=False ) - posterior = jitgp(input_mesh) + posterior = compiled_gp(input_mesh) posterior_mean = posterior.mean.detach().squeeze().cpu().numpy() posterior_std = ( torch.sqrt(posterior.variance).detach().squeeze().cpu().numpy() diff --git a/xopt/tests/generators/bayesian/test_utils.py b/xopt/tests/generators/bayesian/test_utils.py index afe3185f2..657fde9d9 100644 --- a/xopt/tests/generators/bayesian/test_utils.py +++ b/xopt/tests/generators/bayesian/test_utils.py @@ -6,7 +6,6 @@ import pandas as pd import pytest import torch -from botorch.acquisition import UpperConfidenceBound from xopt import Evaluator, Xopt from xopt.generators.bayesian import UpperConfidenceBoundGenerator @@ -15,8 +14,6 @@ compute_hypervolume_and_pf, torch_compile_acqf, torch_compile_gp_model, - torch_trace_acqf, - torch_trace_gp_model, interpolate_points, validate_turbo_controller_base, ) @@ -96,66 +93,6 @@ def test_compute_hypervolume_and_pf(self): assert pf_Y.shape[1] == 3 assert hv > 0 - @pytest.mark.parametrize("use_cuda", cuda_combinations) - def test_model_jit(self, use_cuda): - vocs = deepcopy(TEST_VOCS_BASE) - vocs.constraints = {} - evaluator = Evaluator(function=xtest_callable) - gen = UpperConfidenceBoundGenerator( - vocs=vocs, - ) - gen.use_cuda = use_cuda - X = Xopt(generator=gen, evaluator=evaluator) - gen = X.generator - X.random_evaluate(200) - gen.train_model() - X.random_evaluate(5000) - gen.gp_constructor.use_cached_hyperparameters = True - gen.train_model() - gen.model.eval() - - def get_model(): - return deepcopy(gen.model.models[0]) - - t1 = time.perf_counter() - model_jit = torch_trace_gp_model( - get_model(), gen.vocs, gen.tkwargs, posterior=False, batch_size=500 - ).to(device_map[use_cuda]) - t2 = time.perf_counter() - print(f"JIT compile: {t2 - t1:.4f} seconds") - - t1 = time.perf_counter() - model_jit_posterior = torch_trace_gp_model( - get_model(), gen.vocs, gen.tkwargs, batch_size=500 - ).to(device_map[use_cuda]) - t2 = time.perf_counter() - print(f"JIT posterior compile: {t2 - t1:.4f} seconds") - - x_grid = torch.tensor( - pd.DataFrame( - random_inputs(gen.vocs, 500, include_constants=False) - ).to_numpy() - ) - x_grid = x_grid.to(device_map[use_cuda]) - - m = get_model() - t, values1 = time_call(lambda: m(x_grid), 3) - t = np.array(t) - print(f"Original time: {t[1:].mean():.6f} +- {t[1:].std():.6f}") - - m = get_model() - t, values1 = time_call(lambda: m.posterior(x_grid), 3) - t = np.array(t) - print(f"Original posterior time: {t[1:].mean():.6f} +- {t[1:].std():.6f}") - - t, values1 = time_call(lambda: model_jit(x_grid), 3) - t = np.array(t) - print(f"JIT time: {t[1:].mean():.6f} +- {t[1:].std():.6f}") - - t, values1 = time_call(lambda: model_jit_posterior(x_grid), 3) - t = np.array(t) - print(f"JIT posterior time: {t[1:].mean():.6f} +- {t[1:].std():.6f}") - @pytest.mark.parametrize("use_cuda", cuda_combinations) def test_model_compile(self, use_cuda): # For inductor + windows any, MSVC 2022 build tools are required @@ -202,18 +139,16 @@ def test_model_compile(self, use_cuda): t2 = time.perf_counter() print(f"Compile AT: {t2 - t1:.4f} seconds") - def fmodel(m, x): - mvn = m.posterior(x) - return mvn.mean, mvn.variance - t1 = time.perf_counter() - model_jit = torch_trace_gp_model( - gen.train_model().models[0], - gen.vocs, - gen.tkwargs, + model_compile_direct = torch_compile_gp_model( + gen.train_model().models[0], gen.vocs, gen.tkwargs, posterior=False ).to(device_map[use_cuda]) t2 = time.perf_counter() - print(f"JIT trace: {t2 - t1:.4f} seconds") + print(f"Compile direct: {t2 - t1:.4f} seconds") + + def fmodel(m, x): + mvn = m.posterior(x) + return mvn.mean, mvn.variance x_grid = torch.tensor( pd.DataFrame( @@ -226,44 +161,43 @@ def fmodel(m, x): t1, values1 = time_call(lambda: fmodel(model, x_grid), 10) t1 = np.array(t1) - t2, values2 = time_call(lambda: model_jit(x_grid), 10) + t2, values2 = time_call(lambda: fmodel(model_compile, x_grid), 10) t2 = np.array(t2) - t3, values3 = time_call(lambda: fmodel(model_compile, x_grid), 10) + t3, _ = time_call(lambda: fmodel(model_compile_reduce_overhead, x_grid), 10) t3 = np.array(t3) - t4, values4 = time_call( - lambda: fmodel(model_compile_reduce_overhead, x_grid), 10 - ) + t4, _ = time_call(lambda: fmodel(model_compile_max_autotune, x_grid), 10) t4 = np.array(t4) - t5, values5 = time_call( - lambda: fmodel(model_compile_max_autotune, x_grid), 10 - ) - t5 = np.array(t5) - print(f"Original time: {t1} seconds") - print(f"JIT time: {t2} seconds") - print(f"Compiled time: {t3} seconds") - print(f"Compiled RO time: {t4} seconds") - print(f"Compiled AT time: {t5} seconds") + print(f"Compiled time: {t2} seconds") + print(f"Compiled RO time: {t3} seconds") + print(f"Compiled AT time: {t4} seconds") print(f"Avg: {t1[1:].mean():.6f} +- {t1[1:].std():.6f}") print(f"Avg: {t2[1:].mean():.6f} +- {t2[1:].std():.6f}") print(f"Avg: {t3[1:].mean():.6f} +- {t3[1:].std():.6f}") print(f"Avg: {t4[1:].mean():.6f} +- {t4[1:].std():.6f}") - print(f"Avg: {t5[1:].mean():.6f} +- {t5[1:].std():.6f}") - for v1, v2, v3 in zip(values1, values2, values3): + for v1, v2 in zip(values1, values2): m1, var1 = v1 m2, var2 = v2 - m3, var3 = v3 - assert torch.allclose(m1, m2, rtol=0), "JIT model output mismatch" - assert torch.allclose(var1, var2, rtol=0), "JIT model variance mismatch" - assert torch.allclose(m1, m3, rtol=0), "Compiled model output mismatch" - assert torch.allclose(var1, var3, rtol=0), ( + assert torch.allclose(m1, m2, rtol=0), "Compiled model output mismatch" + assert torch.allclose(var1, var2, rtol=0), ( "Compiled model variance mismatch" ) + with torch.no_grad(), gpytorch.settings.fast_pred_var(): + model.eval() + mvn_eager = model(x_grid) + mvn_direct = model_compile_direct(x_grid) + assert torch.allclose(mvn_eager.mean, mvn_direct.mean, rtol=0), ( + "Compiled model direct call output mismatch" + ) + assert torch.allclose(mvn_eager.variance, mvn_direct.variance, rtol=0), ( + "Compiled model direct call variance mismatch" + ) + @pytest.mark.parametrize("use_cuda", cuda_combinations) def test_acqf_compile(self, use_cuda): print(f"{torch._dynamo.list_backends()=}") @@ -319,14 +253,6 @@ def make_acqf(): t2 = time.perf_counter() print(f"Compile AT: {t2 - t1:.4f} seconds") - acqf = make_acqf().to(device_map[use_cuda]) - t1 = time.perf_counter() - model_jit = torch_trace_acqf(acqf, gen.vocs, gen.tkwargs).to( - device_map[use_cuda] - ) - t2 = time.perf_counter() - print(f"JIT trace: {t2 - t1:.4f} seconds") - def fmodel(m, x): return m(x) @@ -340,47 +266,28 @@ def fmodel(m, x): t1, values1 = time_call(lambda: fmodel(model, x_grid), 10) t1 = np.array(t1) - t2, values2 = time_call(lambda: fmodel(model_jit, x_grid), 10) + t2, values2 = time_call(lambda: fmodel(model_compile, x_grid), 10) t2 = np.array(t2) - t3, values3 = time_call(lambda: fmodel(model_compile, x_grid), 10) + t3, _ = time_call(lambda: fmodel(model_compile_reduce_overhead, x_grid), 10) t3 = np.array(t3) - t4, values4 = time_call( - lambda: fmodel(model_compile_reduce_overhead, x_grid), 10 - ) + t4, _ = time_call(lambda: fmodel(model_compile_max_autotune, x_grid), 10) t4 = np.array(t4) - t5, values5 = time_call( - lambda: fmodel(model_compile_max_autotune, x_grid), 10 - ) - t5 = np.array(t5) - print(f"Original time: {t1} seconds") - print(f"JIT time: {t2} seconds") - print(f"Compiled time: {t3} seconds") - print(f"Compiled RO time: {t4} seconds") - print(f"Compiled AT time: {t5} seconds") + print(f"Compiled time: {t2} seconds") + print(f"Compiled RO time: {t3} seconds") + print(f"Compiled AT time: {t4} seconds") print(f"Original Avg: {t1[1:].mean():.6f} +- {t1[1:].std():.6f}") - print(f"JIT Avg: {t2[1:].mean():.6f} +- {t2[1:].std():.6f}") - print(f"Compiled Avg: {t3[1:].mean():.6f} +- {t3[1:].std():.6f}") - print(f"Compiled RO Avg: {t4[1:].mean():.6f} +- {t4[1:].std():.6f}") - print(f"Compiled AT Avg: {t5[1:].mean():.6f} +- {t5[1:].std():.6f}") + print(f"Compiled Avg: {t2[1:].mean():.6f} +- {t2[1:].std():.6f}") + print(f"Compiled RO Avg: {t3[1:].mean():.6f} +- {t3[1:].std():.6f}") + print(f"Compiled AT Avg: {t4[1:].mean():.6f} +- {t4[1:].std():.6f}") - for v1, v2, v3 in zip(values1, values2, values3): + for v1, v2 in zip(values1, values2): m1 = v1 m2 = v2 - m3 = v3 assert torch.allclose(m1, m2, rtol=1e-5) - assert torch.allclose(m1, m3, rtol=1e-5) - - def test_trace_gp_model_model_list_error(self): - from botorch.models import ModelListGP - - vocs = deepcopy(TEST_VOCS_BASE) - model = ModelListGP() - with pytest.raises(ValueError): - torch_trace_gp_model(model, vocs, {}, posterior=True) def test_compile_gp_model_model_list_error(self): from botorch.models import ModelListGP @@ -390,34 +297,6 @@ def test_compile_gp_model_model_list_error(self): with pytest.raises(ValueError): torch_compile_gp_model(model, vocs, {}, posterior=True) - def test_torch_trace_acqf(self): - evaluator = Evaluator(function=xtest_callable) - gen = UpperConfidenceBoundGenerator( - vocs=TEST_VOCS_BASE, - ) - gen.numerical_optimizer.n_restarts = 2 - gen.n_monte_carlo_samples = 4 - X = Xopt(generator=gen, evaluator=evaluator) - X.random_evaluate(100) - for _ in range(1): - X.step() - - gen = X.generator - model = gen.train_model().models[0] - acq = UpperConfidenceBound(model, beta=0.1) - tkwargs = {"device": torch.device("cpu"), "dtype": torch.double} - traced_acq = torch_trace_acqf(acq, gen.vocs, tkwargs) - assert isinstance(traced_acq, torch.jit.ScriptModule) - # Check output shape matches original - rand_point = random_inputs(gen.vocs)[0] - rand_vec = torch.stack( - [rand_point[k] * torch.ones(1) for k in gen.vocs.variable_names], dim=1 - ).to(**tkwargs) - test_x = rand_vec.unsqueeze(-2) - orig_out = acq(test_x) - traced_out = traced_acq(test_x) - assert torch.allclose(orig_out, traced_out, rtol=1e-6) - def test_interpolate_points_invalid_rows(self): # Create a DataFrame with more than two rows df_invalid = pd.DataFrame({"x": [0, 1, 2], "y": [0, 1, 2]}) diff --git a/xopt/tests/generators/bayesian/test_visualize.py b/xopt/tests/generators/bayesian/test_visualize.py index a181cf7d7..b2783db8e 100644 --- a/xopt/tests/generators/bayesian/test_visualize.py +++ b/xopt/tests/generators/bayesian/test_visualize.py @@ -1,3 +1,4 @@ +import time from unittest.mock import MagicMock import numpy as np @@ -121,6 +122,31 @@ def test_visualize_model(vocs, data, variable_names, show_acquisition): ) +@pytest.mark.parametrize("variable_names", [["x"], ["x", "y"]]) +def test_visualize_model_compile_inductor(vocs, data, variable_names): + generator = UpperConfidenceBoundGenerator(vocs=vocs) + generator.add_data(data) + generator.train_model() + + t1 = time.perf_counter() + fig, ax = visualize.visualize_model( + model=generator.model, + vocs=vocs, + data=data, + tkwargs={}, + output_names=["z"], + variable_names=variable_names, + n_grid=5, + show_acquisition=False, + model_compile_mode="inductor", + ) + t2 = time.perf_counter() + print( + f"visualize_model ({len(variable_names)}D) with inductor compile: {t2 - t1:.2f} s" + ) + assert ax is not None + + def test_visualize_model_reference_title_visibility(vocs, data): generator = UpperConfidenceBoundGenerator(vocs=vocs) generator.add_data(data)