Skip to content

Remove PyTorch JIT trace functionality #457

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nikitakuklev wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Remove PyTorch JIT trace functionality #457

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a0d5a8a
Remove PyTorch JIT functionality
nikitakuklev Aug 20, 2026
42a04e0
Fix model_compile_mode ignored in 1D visualization
nikitakuklev Aug 20, 2026
76900c7
Add tests for uncovered torch.compile paths
nikitakuklev Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions docs/examples/single_objective_bayes_opt/fast_model_eval.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down
147 changes: 12 additions & 135 deletions xopt/generators/bayesian/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down
37 changes: 15 additions & 22 deletions xopt/generators/bayesian/visualize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
-------
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down
Loading
Loading