Initialize MTGP on GPU

It looks like get_MTGP initilizes the model on cpu by default and it doesn’t have a device argument like the other modelbridges.

https://github.com/facebook/Ax/blob/57ba8714902ac218eb87dc2f90090678aa307a43/ax/modelbridge/factory.py#L340-L344

Is there another way to run the multi-task model with EI on GPU? Thanks a lot!

[danielcohenlive]

I didn’t realize there were other factory functions taking a device arg. Made a PR. Also feel free to submit PRs for this kind of thing in the future.

[soerenjalas]

Thanks for the quick response. I just tried this and ran into some issues. For context, I’m basically trying to reproduce the multi-task optimization tutorial (https://ax.dev/tutorials/multi_task.html). This means that I need a multi-type model, correct? I couldn’t find something like Models.MT_MTGP, or is there yet another workaround?

What now ended up working for me was to adapt get_MTGP to pass the device argument through like the other factory functions do

def get_MTGP_gpu(
    experiment: Experiment,
    data: Data,
    search_space: Optional[SearchSpace] = None,
    dtype: torch.dtype = torch.double,
    device: torch.device = DEFAULT_TORCH_DEVICE,
    trial_index: Optional[int] = None,
) -> TorchModelBridge:
    """Instantiates a Multi-task Gaussian Process (MTGP) model that generates
    points with EI.
    If the input experiment is a MultiTypeExperiment then a
    Multi-type Multi-task GP model will be instantiated.
    Otherwise, the model will be a Single-type Multi-task GP.
    """

    if isinstance(experiment, MultiTypeExperiment):
        trial_index_to_type = {
            t.index: t.trial_type for t in experiment.trials.values()
        }
        transforms = MT_MTGP_trans
        transform_configs = {
            "TrialAsTask": {"trial_level_map": {"trial_type": trial_index_to_type}},
            "ConvertMetricNames": tconfig_from_mt_experiment(experiment),
        }
    else:
        # Set transforms for a Single-type MTGP model.
        transforms = ST_MTGP_trans
        transform_configs = None

    # Choose the status quo features for the experiment from the selected trial.
    # If trial_index is None, we will look for a status quo from the last
    # experiment trial to use as a status quo for the experiment.
    if trial_index is None:
        trial_index = len(experiment.trials) - 1
    elif trial_index >= len(experiment.trials):
        raise ValueError("trial_index is bigger than the number of experiment trials")

    # pyre-fixme[16]: `ax.core.base_trial.BaseTrial` has no attribute `status_quo`.
    status_quo = experiment.trials[trial_index].status_quo
    if status_quo is None:
        status_quo_features = None
    else:
        status_quo_features = ObservationFeatures(
            parameters=status_quo.parameters,
            # pyre-fixme[6]: Expected `Optional[numpy.int64]` for 2nd param but got
            #  `int`.
            trial_index=trial_index,
        )

    return TorchModelBridge(
        experiment=experiment,
        search_space=search_space or experiment.search_space,
        data=data,
        model=BotorchModel(),
        transforms=transforms,
        # pyre-fixme[6]: Expected `Optional[Dict[str, Dict[str,
        #  typing.Union[botorch.acquisition.acquisition.AcquisitionFunction, float,
        #  int, str]]]]` for 6th param but got `Optional[Dict[str,
        #  typing.Union[Dict[str, Dict[str, Dict[int, Optional[str]]]], Dict[str,
        #  typing.Union[botorch.acquisition.acquisition.AcquisitionFunction, float,
        #  int, str]]]]]`.
        transform_configs=transform_configs,
        torch_dtype=torch.double,
        torch_device=device,
        status_quo_features=status_quo_features,
    )