[GENERAL SUPPORT]: Slow MultiObjective Small Space Parallel Get Trials

[chanansh]

Question

I have a budget of batch_size=150 parallel runs on the blackbox.
I have setup a strategy to have SOBOL with batch_size followed by BOTORCH
The generation of the N trials by ax.get_next_trials is so slow, making the optimisation useless.
What am I doing wrong?
The objective is MO with 2 scalars (no noise).
The feature space is very small (see below just 2500 possible values).
It takes about 5-10 second per trial on M3 macbook pro (only one cpu is used).

Please provide any relevant code snippet if applicable.

search_space = SearchSpace(
    parameters=[
        # This is a discrete choice of 34 integer values, sorted in ascending order.
        ChoiceParameter(
            name="parameter_A",
            values=[128, 256, 384, 512, 640, 768, 896, 1024, 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048, 2176, 2304, 2432, 2560, 2688, 2816, 2944, 3072, 3200, 3328, 3456, 3584, 3712, 3840, 3968, 4096, 6144, 8192],
            is_ordered=True,
            value_type="int"
        ),
        # This represents a search over powers of 2, from 2^1 to 2^6.
        RangeParameter(
            name="parameter_B",
            lower=1,
            upper=6,
            value_type="int",
        ),
        # This represents 2^0.
        ChoiceParameter(
            name="parameter_C",
            values=[0],
            is_ordered=True,
            value_type="int"
        ),
        # This represents a search over powers of 2, from 2^0 to 2^5.
        RangeParameter(
            name="parameter_D",
            lower=0,
            upper=5,
            value_type="int",
        ),
        # This represents a search over powers of 2, from 2^0 to 2^3.
        RangeParameter(
            name="parameter_E",
            lower=0,
            upper=3,
            value_type="int",
        ),
        # This represents a search over powers of 2, from 2^0 to 2^6.
        RangeParameter(
            name="parameter_F",
            lower=0,
            upper=6,
            value_type="int",
        ),
        # This represents a search over powers of 2, from 2^0 to 2^5.
        RangeParameter(
            name="parameter_G",
            lower=0,
            upper=5,
            value_type="int",
        )
    ]
)

linear_constraints=[
    (["parameter_D", "parameter_E", "parameter_F", "parameter_G"], "==", 6)
]


# generation strategy

        gs = GenerationStrategy(steps=[
            GenerationStep(model=Generators.SOBOL, num_trials=sobol_trials, max_parallelism=batch_size),
            GenerationStep(
                model=Generators.BOTORCH_MODULAR, 
                num_trials=-1, 
                max_parallelism=batch_size
            ),
        ])

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[Balandat]

This is a very high parallelism setting, the Hypervolume-based multi-objective algorithms we use by default in the BOTORCH_MODULAR strategy are somewhat computationally intensive, so I'm not surprised this takes a long time.

With this batch size, it could make sense to use a faster strategy, such as qParEGO (see this botorch tutorial for more details on this).

To use this, you need to configure your BOTORCH_MODULAR step as follows:

from botorch.acquisition.multi_objective.parego import qLogNParEGO

GenerationStep(
    model=Generators.BOTORCH_MODULAR, 
    model_kwargs={"botorch_acqf_class": qLogNParEGO},
    num_trials=-1, 
    max_parallelism=batch_size,
),

It takes about 5-10 second per trial on M3 macbook pro (only one cpu is used).

So 5-10 seconds to generate 150 points? Or to generate one point?

Taking a step back: How many batches are you able to run? If the total cardinality of the search space is 2500 and you can run 150 evaluations per batch, you could fully enumerate the search space in 17 batches. So I assume you can run fewer than that? How long does the evaluation of one batch take?

linear_constraints=[(["parameter_D", "parameter_E", "parameter_F", "parameter_G"], "==", 6)]

This doesn't seem to be used, but just wanted to point out that linear equality constraints are currently not supported (only inequality constraints).

[chanansh]

Hi,

• 10 second for a single point.
• linear constraint is a typo, I meant inequality.
• I meant to run 2-3 batches.

[chanansh]

@Balandat - would not a for loop for suggestions (aka get_trial) will result in the same suggestion over and over? if not, why? Isn't it just looking for max on the acq fun?

[Balandat]

would not a for loop for suggestions (aka get_trial) will result in the same suggestion over and over? if not, why? Isn't it just looking for max on the acq fun?

It would not; the API will add the features of trials in CANDIDATE stage (which new trials are in by default) as "pending points" (see here), and by doing so the acquisition function can take these points into account when generating the next suggestion. In other words, with pending points the acquisition function will have a different maximum then when first called.

[chanansh]

@Balandat I think I found an issue, please correct me if I am wrong:

ax.get_next_trials is naive and just does a for-loop (albeit you said it actually keeps a state of candidates):

        for _ in range(max_trials):
            try:
                params, trial_index = self.get_next_trial(
                    ttl_seconds=ttl_seconds, fixed_features=fixed_features
                )
                trials_dict[trial_index] = params
            except OptimizationComplete as err:
                logger.info(
                    f"Encountered exception indicating optimization completion: {err}"
                )
                return trials_dict, True

However the qEHVI family of algorithm gets q=batch_size as an input to do it in a "smart" way.
For example, in the BoTorch docs:

    # optimize
    candidates, _ = optimize_acqf(
        acq_function=acq_func,
        bounds=standard_bounds,
        q=BATCH_SIZE,
        num_restarts=NUM_RESTARTS,
        raw_samples=RAW_SAMPLES,  # used for intialization heuristic
        options={"batch_limit": 5, "maxiter": 200},
        sequential=True,
    )

Notice how q is sent to get q candidates.

Could this be the issue with batch MOO in AxDev?

[Balandat]

@chanansh which version of Ax are you referring to? The current version has an updated API and doesn't define a get_next_trial() method (only a get_next_trials() one).

[chanansh]

ax-platform 1.0.0 Adaptive Experimentation

strange. let me check.

I can see the signature

    def get_next_trial(
        self,
        ttl_seconds: int | None = None,
        force: bool = False,
        fixed_features: FixedFeatures | None = None,
    ) -> tuple[TParameterization, int]:

and I can confirm that version.py contains:

__version__ = version = '1.0.0'
__version_tuple__ = version_tuple = (1, 0, 0)

get_next_trials seems to be just a for loop inside (with some constrains)

[lena-kashtelyan]

I think you are looking at the legacy AxClient, not the new Client API (showcased on Ax website, e.g. here: https://ax.dev/docs/tutorials/quickstart/), @chanansh ?

[chanansh]

I was on 1.0.0 I will check. p.s. there is no updated documentation for MOBO - https://ax.dev/docs/0.5.0/tutorials/multiobjective_optimization/ - if you choose a newer version it does not exist anymore.

[lena-kashtelyan]

Hi @chanansh , sorry we'd lost track of this issue! Multi-objective optimization is built into the Client API natively: please check out this recipe: https://ax.dev/docs/recipes/multi-objective-optimization.