Configure an optimization engine

Choose an optimization engine for Weights & Biases Sweeps. Use the built-in Weights & Biases engine for random, grid, or Bayesian search, or use Optuna or Ax for their search algorithms and customization options. The engine, scheduler, and sweep agents work together:

Component	Role
Optimization engine	Suggests hyperparameter configurations and, depending on the search strategy, uses run results to guide the search.
Sweep scheduler	Requests configurations from the engine, submits them to Weights & Biases, and passes run results back to the engine.
Sweep agent	Retrieves assigned configurations from Weights & Biases and executes your training code.

This page describes how to configure an engine for a scheduler running in your own environment. For complete YAML examples, see Sweep scheduler examples.

Choose an engine

Choose an engine based on the search methods and customization options you need:

Engine	Description
Weights & Biases (built-in)	Random, grid, or Bayesian search configured through sweep configuration keys.
Optuna	An open source hyperparameter optimization framework with configurable samplers and pruners.
Ax	An open source platform for adaptive experimentation with configurable generation and stopping strategies.

Use Optuna or Ax to:

Define multiple objectives to optimize more than one metric.
Define a custom sweep-level stopping function to control when the scheduler stops proposing new runs.

With Optuna, you can also define a conditional search space to select parameters based on earlier choices in a trial.

Install dependencies

Install Weights & Biases Python SDK version 0.31.0 or later with the extra dependencies for your engine. Run the installation command in the environment where you run the scheduler.

Weights & Biases
Optuna
Ax

Use Python 3.10 or later. Install the SDK with the sweeps extra to include sweeps>=0.2.0 in the scheduler environment:

pip install "wandb[sweeps]>=0.31.0"

Use Python 3.10 or later. Install the SDK with the optuna extra to include optuna>=4.0.0 in the scheduler environment:

pip install "wandb[optuna]>=0.31.0"

To use the OptunaHub Terminator example on this page, also install optunahub.

pip install optunahub

Use Python 3.11 or later. Install the SDK with the ax extra to include ax-platform>=1.3.1 in the scheduler environment:

pip install "wandb[ax]>=0.31.0"

Configure optimization

Configure the search strategy, objectives, and search space for your engine. The built-in Weights & Biases engine uses the sweep configuration. For Optuna and Ax, you can provide a Python function that configures an Optuna study or Ax client.

Configure the search strategy

Set scheduler.engine to wandb, optuna, or ax. For Optuna and Ax, custom optimizer functions are optional. You can switch engines by changing scheduler.engine and use the search space and objectives defined in the sweep configuration, including multiple objectives defined with metrics. To use a custom optimizer function, set scheduler.source to your Python filename and scheduler.optimizer to the function name. For Optuna, you can also define a custom search space with scheduler.search_space.

Weights & Biases
Optuna
Ax

Configure the built-in Weights & Biases engine using a sweep configuration.Set method to random, grid, or bayes. The engine uses method, metric, and parameters to select configurations for sweep runs.For example, to configure Bayesian search using the built-in Weights & Biases engine:

sweep.yaml

program: train.py
method: bayes
metric:
  name: validation_loss
  goal: minimize
parameters:
  learning_rate:
    min: 0.0001
    max: 0.01
  batch_size:
    values: [16, 32, 64]
scheduler:
  engine: wandb

For details about each search method, see Search methods.

To customize the optimizer, define a function that accepts one SweepInfo argument and returns an Optuna Study.For example, suppose you have the following sweep configuration:

sweep.yaml

program: train.py

# Required by the W&B sweep schema, but the Optuna scheduler
# controls the search strategy.
method: random

metric:
  name: val_loss
  goal: minimize

parameters:
  n_estimators:
    min: 20
    max: 200
  max_depth:
    min: 2
    max: 20
  min_samples_split:
    min: 0.01
    max: 0.5

# Stop after at most 100 runs, even if the terminator does not
# detect convergence first.
run_cap: 100

scheduler:
  engine: optuna
  source: optimizer.py
  optimizer: create_optimizer

In optimizer.py, create a study with a TPE sampler and return it with the Terminator’s stopping function:

optimizer.py

import optuna
import optunahub

terminator_module = optunahub.load_module("callbacks/terminator")

EMMREvaluator = terminator_module.EMMREvaluator
MedianErrorEvaluator = terminator_module.MedianErrorEvaluator
Terminator = terminator_module.Terminator


def create_optimizer(sweep_info):
    study = optuna.create_study(
        direction="minimize",
        sampler=optuna.samplers.TPESampler(seed=42),
    )

    improvement_evaluator = EMMREvaluator()
    error_evaluator = MedianErrorEvaluator(improvement_evaluator)

    terminator = Terminator(
        improvement_evaluator=improvement_evaluator,
        error_evaluator=error_evaluator,
        min_n_trials=20,
    )

    return study, terminator.should_terminate

The previous example configures two parts of the Optuna integration:

A TPESampler, which suggests hyperparameter configurations.
An OptunaHub Terminator, which stops the sweep when further improvement is unlikely.

Don’t call study.optimize() in your optimizer function. The scheduler uses Optuna’s ask() and tell() methods to request trials and report results. This integration doesn’t expose the Study.optimize() callback interface.

To customize the optimizer, define a function that accepts one SweepInfo argument and returns an Ax Client.Import Ax helper functions to configure the client’s search space and objectives from the sweep configuration:

Helper function	Description
`sweep_config_to_search_space(config)`	Converts `parameters` into the parameter configurations accepted by `Client.configure_experiment(parameters=...)`.
`configure_sweep_objective(client, config)`	Configures optimization from `metric` or `metrics`. It preserves metric names containing characters such as hyphens, which Ax’s expression parser can misinterpret. We recommend this helper over `Client.configure_optimization()` for these names.

Mismatched objective names or directions fail validation when the scheduler starts.For example, the following function creates an Ax client and configures it from sweep_info.config:

optimizer.py

from ax.api.client import Client
from wandb.sdk.sweeps.scheduler.ax import (
    configure_sweep_objective,
    sweep_config_to_search_space,
)

def create_optimizer(sweep_info):
    client = Client()
    client.configure_experiment(
        parameters=sweep_config_to_search_space(sweep_info.config)
    )
    configure_sweep_objective(client, sweep_info.config)
    return client

In sweep.yaml, set scheduler.source to optimizer.py and scheduler.optimizer to create_optimizer:

sweep.yaml

program: train.py
method: custom
metric:
  name: validation_loss
  goal: minimize
parameters:
  learning_rate:
    min: 0.0001
    max: 0.01
  batch_size:
    values: [16, 32, 64]
scheduler:
  engine: ax
  source: optimizer.py
  optimizer: create_optimizer

For individual-run and sweep-level stopping options, see Configure stopping.

Define multiple objectives

If you use Optuna or Ax, you can use metrics (plural) to define multiple objectives. When you create a sweep with wandb sweep, Weights & Biases stores the sweep configuration. When the scheduler starts, it retrieves that configuration and passes it to your custom optimizer function as part of a SweepInfo object. Use sweep_info.config to access the stored sweep configuration.

The built-in Weights & Biases engine supports only a single objective, defined with metric.

You don’t need a custom optimizer function to use multiple objectives. The following examples show how to configure the objectives when you provide one.

Optuna
Ax

For a multi-objective sweep with Optuna, define each objective in the metrics list. Use the sweep_directions() helper from wandb.sdk.sweeps.scheduler.optuna to convert the goals to Optuna study directions in the same order.For example, the following sweep minimizes validation loss and maximizes accuracy:

sweep.yaml

name: optuna-multi-objective
program: train.py

scheduler:
  engine: optuna
  source: optimizer.py
  optimizer: create_optimizer

metrics:
  - name: validation_loss
    goal: minimize
  - name: accuracy
    goal: maximize

parameters:
  learning_rate:
    min: 0.0001
    max: 0.01
  batch_size:
    values: [16, 32, 64]

The following demo script reads the suggested learning_rate and batch_size values, calculates two synthetic objective values, and logs them to the run. The final logged values appear in the run summary, where the scheduler can retrieve them.

train.py

import math
import wandb

with wandb.init() as run:
    learning_rate = run.config["learning_rate"]
    batch_size = run.config["batch_size"]

    # Replace these calculations with model training and evaluation.
    validation_loss = (
        abs(math.log10(learning_rate) + 3)
        + 16 / batch_size
    )
    accuracy = min(
        1.0,
        0.8 - abs(math.log10(learning_rate) + 2.5) * 0.05 + batch_size / 1000,
    )

    run.log(
        {
            "validation_loss": validation_loss,
            "accuracy": accuracy,
        }
    )

Within your optimizer.py file, pass the stored sweep configuration from sweep_info.config to sweep_directions() when you create the Optuna study:

optimizer.py

import optuna
from wandb.sdk.sweeps.scheduler.optuna import sweep_directions

def create_optimizer(sweep_info):
    return optuna.create_study(
        directions=sweep_directions(sweep_info.config),
    )

The resulting study minimizes validation_loss and maximizes accuracy.

For a multi-objective sweep with Ax, use the configure_sweep_objective() helper from wandb.sdk.sweeps.scheduler.ax to configure the client with the name and optimization direction of each entry in metrics.For example, the following sweep minimizes validation loss and maximizes accuracy:

sweep.yaml

name: ax-multi-objective
program: train.py

scheduler:
  engine: ax
  source: optimizer.py
  optimizer: create_optimizer

metrics:
  - name: validation_loss
    goal: minimize
  - name: accuracy
    goal: maximize

parameters:
  learning_rate:
    min: 0.0001
    max: 0.01
  batch_size:
    values: [16, 32, 64]

train.py

import math
import wandb

with wandb.init() as run:
    learning_rate = run.config["learning_rate"]
    batch_size = run.config["batch_size"]

    # Replace these calculations with model training and evaluation.
    validation_loss = (
        abs(math.log10(learning_rate) + 3)
        + 16 / batch_size
    )
    accuracy = min(
        1.0,
        0.8 - abs(math.log10(learning_rate) + 2.5) * 0.05 + batch_size / 1000,
    )

    run.log(
        {
            "validation_loss": validation_loss,
            "accuracy": accuracy,
        }
    )

In optimizer.py, create an Ax client and configure its search space and objectives from the stored sweep configuration. Use sweep_config_to_search_space() to convert parameters into Ax parameter definitions, and configure_sweep_objective() to configure the objectives defined in metrics.

optimizer.py

from ax.api.client import Client
from wandb.sdk.sweeps.scheduler.ax import (
    configure_sweep_objective,
    sweep_config_to_search_space,
)


def create_optimizer(sweep_info):
    client = Client()
    client.configure_experiment(
        parameters=sweep_config_to_search_space(sweep_info.config)
    )
    configure_sweep_objective(client, sweep_info.config)
    return client

In this example, configure_sweep_objective() configures Ax to minimize validation_loss and maximize accuracy.

Define a conditional search space

Only the Optuna engine supports scheduler.search_space.

For Optuna, use scheduler.search_space to define conditional parameters in Python. For example, you can define different hyperparameters for different model types or vary the number of parameters based on earlier choices. Set scheduler.search_space to the name of a function in the file specified by scheduler.source. The scheduler passes an Optuna trial to the function and uses the returned dictionary as the run configuration.

For a fixed search space, define parameters in the sweep configuration.

When you set scheduler.search_space, the scheduler ignores the top-level parameters definitions. Define every parameter you want to pass to a run in the Python function. For example, suppose you have the following sweep configuration:

sweep.yaml

metric:
  name: validation_loss
  goal: minimize

parameters: {}

scheduler:
  engine: optuna
  source: search.py
  search_space: define_search_space

For each trial, the scheduler calls define_search_space(trial) and uses the returned dictionary as the run configuration:

search.py

def define_search_space(trial):
    model = trial.suggest_categorical("model", ["linear", "transformer"])

    config = {
        "model": model,
        "learning_rate": trial.suggest_float(
            "learning_rate", 1e-5, 1e-2, log=True
        ),
    }

    if model == "transformer":
        config["num_layers"] = trial.suggest_int("num_layers", 2, 12)
        config["num_heads"] = trial.suggest_categorical(
            "num_heads", [4, 8, 16]
        )
    else:
        config["regularization"] = trial.suggest_float(
            "regularization", 0.0, 1.0
        )

    return config

Every run receives model and learning_rate. Runs with model="transformer" also receive num_layers and num_heads. Runs with model="linear" receive regularization instead.

Configure stopping

Stopping can apply to an individual run or to the overall sweep. Individual-run early stopping ends runs that are unlikely to produce a useful result. Sweep-level stopping prevents the scheduler from proposing new runs after a condition is met.

Stop individual runs

To evaluate a run’s progress, log intermediate values for the objective metric throughout training. Each engine provides its own mechanism for deciding whether to stop a run.

The top-level early_terminate key applies only to the built-in Weights & Biases engine. For Optuna or Ax, configure the engine’s native early-stopping mechanism.

Weights & Biases
Optuna
Ax

Use Hyperband for hyperparameter optimization with early stopping based on intermediate results.Set early_terminate.type to hyperband and use min_iter to set the minimum number of iterations.For example, the following sweep configuration uses Hyperband for early stopping with a minimum of 3 iterations:

sweep.yaml

program: train.py
method: bayes
metric:
  name: validation_loss
  goal: minimize
parameters:
  learning_rate:
    min: 0.0001
    max: 0.01
  batch_size:
    values: [16, 32, 64]
early_terminate:
  type: hyperband
  min_iter: 3
scheduler:
  engine: wandb

See early_terminate for configuration options.

In Optuna, a pruner decides whether to stop a trial early based on intermediate results.For example, suppose you have the following sweep configuration and want to apply early stopping with Optuna:

sweep.yaml

name: optuna-pruning-demo
program: train.py
method: grid
run_cap: 2

metric:
  name: loss
  goal: minimize

parameters:
  loss:
    values: [0.2, 0.9]

scheduler:
  engine: optuna
  source: optimizer.py
  optimizer: make_study

In optimizer.py, create a study with a grid sampler and a threshold pruner. The pruner recommends stopping trials whose intermediate loss exceeds 0.5:

optimizer.py

import optuna

def make_study(sweep_info):
    """Create a study with a grid sampler and a threshold pruner."""
    return optuna.create_study(
        direction="minimize",
        sampler=optuna.samplers.GridSampler(
            {"loss": [0.2, 0.9]},
        ),
        # Stop any run whose intermediate loss exceeds 0.5.
        pruner=optuna.pruners.ThresholdPruner(
            upper=0.5,
            n_warmup_steps=0,
        ),
    )

In train.py, log the configured loss at each step so the scheduler can pass intermediate results to the pruner:

train.py

import time
import wandb

with wandb.init() as run:
    loss = run.config["loss"]

    for step in range(30):
        run.log({"loss": loss}, step=step)
        time.sleep(1)

Optuna pruning is not available for multi-objective studies.

To enable individual-run early stopping with Ax, call client.set_early_stopping_strategy() to configure an early-stopping strategy, then return the client.For example, suppose you have the following sweep configuration and want to apply early stopping with Ax:

sweep.yaml

name: ax-run-stopping-demo
program: train.py

# Required by the sweep schema; Ax controls parameter generation.
method: random
run_cap: 2
metric:
  name: loss
  goal: minimize
parameters:
  difficulty:
    values: [0.1, 1.0]
scheduler:
  engine: ax
  source: optimizer.py
  optimizer: make_client

In optimizer.py, configure the client with the search space, objective, and a percentile early-stopping strategy. Because the objective minimizes loss, the strategy uses higher loss values to identify underperforming runs.

optimizer.py

from ax.api.client import Client
from ax.early_stopping.strategies import PercentileEarlyStoppingStrategy

from wandb.sdk.sweeps.scheduler.ax import (
    configure_sweep_objective,
    sweep_config_to_search_space,
)

def make_client(sweep_info):
    client = Client()
    client.configure_experiment(
        parameters=sweep_config_to_search_space(sweep_info.config),
    )
    configure_sweep_objective(client, sweep_info.config)

    client.set_early_stopping_strategy(
        PercentileEarlyStoppingStrategy(
            percentile_threshold=50.0,
            min_progression=2,
            min_curves=2,
            seconds_between_polls=0,
            normalize_progressions=False,
        )
    )

    return client

In train.py, calculate a synthetic loss from difficulty and log it at each step. These intermediate values let Ax compare the runs’ progress:

train.py

import time
import wandb

with wandb.init() as run:
    difficulty = run.config["difficulty"]

    for step in range(30):
        # The run with difficulty=1.0 is consistently worse.
        loss = difficulty + 1.0 / (step + 1)

        run.log(
            {"loss": loss},
            step=step,
        )

        # Keep the runs alive long enough for the scheduler to poll them.
        time.sleep(1)

Run two sweep agents concurrently so the scheduler can collect intermediate results from both runs. When Ax recommends stopping a run, the scheduler requests that Weights & Biases stop it.

Stop the sweep

Sweep-level stopping controls when the scheduler stops proposing new runs. Use run_cap to set a fixed maximum number of runs for any engine. For Optuna or Ax, you can also return a stopping function that evaluates the current state of the optimization.

In the following examples, the scheduler controls the search strategy, so it ignores method: random. However, the Weights & Biases sweep schema requires method.

The scheduler might already have runs queued or running when a stopping condition is met. Those runs can continue after the scheduler stops proposing new ones.

Optuna
Ax

To enable sweep-level stopping, the optimization function you specify in scheduler.optimizer must return an Optuna study and a stopping function as a tuple:

(study, terminator_callback)

After processing results, the scheduler passes the current study to the stopping function. The function returns True to stop the sweep or False to continue proposing new runs.For example, the following sweep configuration searches for x, limits the sweep to 20 runs, and instructs the Optuna scheduler to call make_study() from the optimizer.py file:

sweep.yaml

name: optuna-sweep-stopping-demo
program: train.py
method: random
run_cap: 20

metric:
  name: loss
  goal: minimize

parameters:
  x:
    min: 0
    max: 10

scheduler:
  engine: optuna
  source: optimizer.py
  optimizer: make_study

In this example, run_cap provides a fallback limit if the stopping function doesn’t end the sweep first.Next, define the Optuna study and stopping function. make_study() configures a seeded random sampler and returns the study with stop_after_three_completed_trials():

optimizer.py

import optuna


def stop_after_three_completed_trials(study: optuna.Study) -> bool:
    """Stop the sweep after three completed trials."""
    completed = sum(
        trial.state == optuna.trial.TrialState.COMPLETE
        for trial in study.trials
    )

    return completed >= 3


def make_study(sweep_info):
    study = optuna.create_study(
        direction="minimize",
        sampler=optuna.samplers.RandomSampler(seed=42),
    )

    return study, stop_after_three_completed_trials

The stopping function counts only successfully completed trials. It returns True after the third completed trial, which tells the scheduler to stop proposing configurations. Failed, pruned, and running trials don’t contribute to the count.Finally, define the objective that each sweep run evaluates. The training script retrieves the suggested value of x, calculates its loss, and logs that loss to Weights & Biases:

train.py

import wandb

with wandb.init() as run:
    x = run.config["x"]
    loss = (x - 3) ** 2

    run.log({"loss": loss})

After each run finishes, the scheduler passes its logged loss to the Optuna study. Once three trials complete, the stopping function returns True and the scheduler stops proposing new runs. Runs that are already queued or running can still finish.

To enable sweep-level stopping, return a tuple containing the Ax client and a stopping function from the function specified by scheduler.optimizer:

(client, terminator)

The following configuration limits the sweep to 20 runs. The stopping function can end it earlier if the objective stops improving:

sweep.yaml

name: ax-sweep-stopping-example
program: train.py
method: random
run_cap: 20
metric:
  name: loss
  goal: minimize
parameters:
  x:
    min: 0.0
    max: 1.0
scheduler:
  engine: ax
  source: optimizer.py
  optimizer: make_client

In optimizer.py, create an Ax client and a stopping function. This example uses ImprovementGlobalStoppingStrategy to compare recent improvement with the spread of observed objective values:

optimizer.py

from ax.api.client import Client
from ax.global_stopping.strategies.improvement import (
    ImprovementGlobalStoppingStrategy,
)

from wandb.sdk.sweeps.scheduler.ax import (
    configure_sweep_objective,
    sweep_config_to_search_space,
)


def make_client(sweep_info):
    """Create the Ax client and sweep-level stopping callback."""

    client = Client()

    # Convert the parameters in sweep.yaml into an Ax search space.
    client.configure_experiment(
        parameters=sweep_config_to_search_space(sweep_info.config),
    )

    # Configure "loss" as the objective Ax should minimize.
    configure_sweep_objective(client, sweep_info.config)

    client.configure_generation_strategy(
        method="random_search",
        initialization_random_seed=7,
        initialize_with_center=True,
    )

    stopping_strategy = ImprovementGlobalStoppingStrategy(
        min_trials=6,
        window_size=3,
        improvement_bar=0.01,
        inactive_when_pending_trials=True,
    )

    def should_terminate_sweep(client):
        """Return True when Ax recommends ending the whole sweep."""
        should_stop, reason = stopping_strategy.should_stop_optimization(
            experiment=client._experiment,
        )

        if should_stop:
            print(f"Ax is stopping the sweep: {reason}")

        return should_stop

    return client, should_terminate_sweep

The following training script evaluates a synthetic objective whose minimum occurs at x=0.5. The generation strategy evaluates that point first, so subsequent trials cannot improve the best result:

train.py

import time
import wandb

with wandb.init() as run:
    x = run.config["x"]

    for step in range(5):
        loss = (x - 0.5) ** 2 + 1.0 / (step + 1)

        run.log(
            {"loss": loss},
            step=step,
        )

        time.sleep(0.5)

This stopping strategy uses only the final objective value from each completed run. The script logs intermediate values to simulate training; this example does not configure individual-run early stopping.

Handle scheduler restarts

The sweep scheduler operates on an existing sweep, so you can stop and restart the scheduler without creating a new sweep. Each time the scheduler starts, it reconstructs the optimization history from the sweep’s existing run data and then continues the optimization loop. It does not restore the engine from a saved optimizer checkpoint. When the scheduler starts or restarts, it calls your custom optimizer function to create a new engine instance. Return a fresh Optuna study or Ax client each time the function runs. Don’t load a separate optimizer checkpoint or reuse an engine instance from a previous scheduler process. If you use a scheduler.source file, the scheduler loads its current contents each time it starts. You can update the file between starts to fix bugs or modify the custom optimizer or search_space function. These changes can affect subsequent suggestions, so preserve the version of the source file used for each scheduler session if you need to reproduce the optimization.

Weights & Biases

Guides

Integrations

Reference

Support

Configure an optimization engine

Choose an engine

Install dependencies

Configure optimization

Configure the search strategy

Define multiple objectives

Define a conditional search space

Configure stopping

Stop individual runs

Stop the sweep

Handle scheduler restarts

​Choose an engine

​Install dependencies

​Configure optimization

​Configure the search strategy

​Define multiple objectives

​Define a conditional search space

​Configure stopping

​Stop individual runs

​Stop the sweep

​Handle scheduler restarts

Choose an engine

Install dependencies

Configure optimization

Configure the search strategy

Define multiple objectives

Define a conditional search space

Configure stopping

Stop individual runs

Stop the sweep

Handle scheduler restarts