Hyperparameter optimization

A sweep trains the same model many times with different settings and keeps the best one. Most of its cost goes to trials that turn out to be bad. Flyte makes each trial a task, so trials run in parallel on their own hardware, retry on their own, and are cached. The driver that chooses which trials to run is ordinary Python, so you can stop weak trials early and give their budget to strong ones.

Run every trial as a task

Put trials and the driver in separate environments. Each trial gets the hardware one training run needs, and the driver only chooses settings and compares scores:

hpo.py
                
                    
                
            
                
            

                
                    
                
            
import asyncio
import math
import random
import tempfile
from dataclasses import dataclass
from pathlib import Path

import flyte
import joblib
import pandas as pd
from flyte.io import File

image = flyte.Image.from_uv_script(__file__, name="hpo")

trial_env = flyte.TaskEnvironment(
    name="hpo-trial",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)

driver_env = flyte.TaskEnvironment(
    name="hpo-driver",
    image=image,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    depends_on=[trial_env],
)

The trial task trains one setting for a number of epochs and returns its validation score with a checkpoint. It can also continue from a checkpoint, which the next section uses:

hpo.py
                
                    
                
            
                
            

                
                    
                
            
@dataclass
class Params:
    alpha: float
    eta0: float

@dataclass
class Trial:
    params: Params
    epochs: int
    score: float
    checkpoint: File

@trial_env.task(cache="auto", retries=2)
async def train_trial(
    df: pd.DataFrame, params: Params, epochs: int, start_epoch: int = 0, checkpoint: File | None = None
) -> Trial:
    """Train one setting up to `epochs`, continuing from `checkpoint` if one is given."""
    from sklearn.linear_model import SGDClassifier
    from sklearn.model_selection import train_test_split

    X, y = df.drop(columns="target").to_numpy() / 16.0, df["target"].to_numpy()
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=0)

    if checkpoint is not None:
        model = joblib.load(await checkpoint.download())
    else:
        model = SGDClassifier(alpha=params.alpha, eta0=params.eta0, learning_rate="constant", random_state=0)

    for _ in range(start_epoch, epochs):
        model.partial_fit(X_train, y_train, classes=sorted(set(y)))

    path = Path(tempfile.mkdtemp()) / "model.joblib"
    joblib.dump(model, path)
    return Trial(
        params=params,
        epochs=epochs,
        score=float(model.score(X_val, y_val)),
        checkpoint=await File.from_local(str(path)),
    )

cache="auto" stores every finished trial. Rerunning the sweep with more trials runs only the new ones, and a sweep that failed halfway resumes without retraining what finished.
retries=2 reruns a trial that hit a bad node, without failing the sweep.
Typed results. Params and Trial are dataclasses, so the UI shows each trial’s settings, score, and checkpoint.

Sweep in parallel

A random search samples settings and launches every trial at once. A semaphore caps how many run at the same time, so a large sweep doesn’t take the whole cluster:

hpo.py
                
def sample_params(rng: random.Random) -> Params:
    return Params(alpha=10 ** rng.uniform(-6, -2), eta0=10 ** rng.uniform(-3, 0))

@driver_env.task
async def random_search(n_trials: int = 16, max_parallel: int = 8, seed: int = 0) -> Params:
    df = await load_data()
    rng = random.Random(seed)
    limit = asyncio.Semaphore(max_parallel)

    async def run_one(i: int, params: Params) -> Trial:
        async with limit:
            with flyte.group(f"trial-{i}"):
                return await train_trial(df, params, epochs=20)

    trials = await asyncio.gather(*(run_one(i, sample_params(rng)) for i in range(n_trials)))
    return max(trials, key=lambda t: t.score).params

Sample from a seeded random.Random, so a rerun proposes the same settings and gets every finished trial from the cache. Random search usually finds a good setting with fewer trials than a grid, because it doesn’t spend trials on dimensions that don’t matter. For a grid over a few values, see Classical machine learning.

Stop weak trials early

Most trials are bad after a few epochs as well as at the end. Successive halving trains every trial on a small budget, keeps the best fraction, and continues only those from their checkpoints with a larger budget:

hpo.py
                
                    
                
            
                
            

                
                    
                
            
@driver_env.task(report=True)
async def successive_halving(n_trials: int = 16, min_epochs: int = 2, eta: int = 2, seed: int = 0) -> Trial:
    df = await load_data()
    rng = random.Random(seed)
    trials = [None] * n_trials
    params = [sample_params(rng) for _ in range(n_trials)]
    rounds = int(math.log(n_trials, eta))
    history = []

    epochs = min_epochs
    for rung in range(rounds + 1):
        with flyte.group(f"rung-{rung}"):
            trials = await asyncio.gather(
                *(
                    train_trial(
                        df,
                        p,
                        epochs=epochs,
                        start_epoch=t.epochs if t else 0,
                        checkpoint=t.checkpoint if t else None,
                    )
                    for p, t in zip(params, trials)
                )
            )
        history.append((rung, epochs, len(trials), max(t.score for t in trials)))
        if rung == rounds:
            break

        # Keep the best 1/eta of the trials and give each eta times the budget.
        survivors = sorted(trials, key=lambda t: -t.score)[: max(1, len(trials) // eta)]
        params = [t.params for t in survivors]
        trials = survivors
        epochs *= eta

    rows = "".join(f"<tr><td>{r}</td><td>{e}</td><td>{n}</td><td>{s:.4f}</td></tr>" for r, e, n, s in history)
    await flyte.report.replace.aio(
        "<h2>Successive halving</h2><table><tr><th>Rung</th><th>Epochs</th><th>Trials</th>"
        f"<th>Best accuracy</th></tr>{rows}</table>",
        do_flush=True,
    )
    return max(trials, key=lambda t: t.score)

With 16 trials and eta=2, the sweep trains 16 trials for 2 epochs, 8 for 4, 4 for 8, 2 for 16, and 1 for 32. Because survivors continue from their checkpoints, that is 96 epochs of training in total, against 512 to train all 16 for 32 epochs. Each rung is a group in the UI. A surviving trial passes its checkpoint to its next rung, so it continues training instead of starting over. Hyperband runs several of these brackets with different starting budgets, and the same driver loop covers it.

Search adaptively with Optuna

To let earlier trials decide which settings to try next, drive an Optuna study from the driver task. The study asks for a setting, and self.objective is a Flyte task that scores it, so each trial runs as its own action:

optimizer.py
                
    async def spawn(self, semaphore: asyncio.Semaphore):
        async with semaphore:
            trial: Trial = self.study.ask()

            try:
                print("Starting trial", trial.number)

                params = {
                    "n_estimators": trial.suggest_int("n_estimators", 10, 200),
                    "max_depth": trial.suggest_int("max_depth", 2, 20),
                    "min_samples_split": trial.suggest_float(
                        "min_samples_split", 0.1, 1.0
                    ),
                }

                output = await self.objective(params)

                self.study.tell(trial, output, state=optuna.trial.TrialState.COMPLETE)
            except flyte.errors.RuntimeUserError as e:
                print(f"Trial {trial.number} failed: {e}")

                self.study.tell(trial, state=optuna.trial.TrialState.FAIL)

            await asyncio.sleep(self.delay)

A semaphore bounds the trials in flight, and a trial that raises flyte.errors.RuntimeUserError is recorded as failed instead of stopping the study. Optuna’s samplers, such as TPE, propose each new setting from the results so far, so the number of trials in flight trades speed against how much each proposal can learn. See the Hyperparameter optimization tutorial for the full optimizer.

Use a sweep tool you already have

These integrations run their own search and use Flyte to execute the trials:

Integration	How it sweeps	Each trial runs as
Weights & Biases sweeps	`@wandb_sweep` creates a W&B sweep, and parallel agent tasks pull settings from it. W&B chooses the settings with grid, random, or Bayesian search, and plots every trial.	One call in a sweep agent task
Hydra	`--multirun` expands config overrides into a grid, and `hydra/sweeper=optuna` runs an Optuna sweeper over your Hydra config.	A separate Flyte run
Ray	Ray Tune runs its schedulers, such as ASHA and PBT, on a Ray cluster that a task starts.	A Ray trial inside one task

To record every trial’s settings and metrics in an experiment tracker, log from the trial task with Weights & Biases or MLflow.

Size trials for the hardware

GPU trials. Give each trial one GPU, and cap concurrency at the number of GPUs you want the sweep to use. A trial that runs out of memory raises flyte.errors.OOMError to the driver, which can record it as failed. See Durable training.
Multi-node trials. Make the trial task a flyte.clustered.ClusteredTaskEnvironment. The driver doesn’t change. See Pre-training.
Preemptible capacity. Trials are short and retry on their own, so they suit interruptible=True. Keep the driver on on-demand capacity. See Use spot capacity.

Seconds-long trials. When a trial takes less time than a container takes to start, use reusable containers to send trials to a warm pool of replicas.

Hyperparameter optimization: a parallel Optuna study, end to end.
Controlling parallel execution: limit concurrency for a whole run.
Evaluation: score the winning setting on a held-out set before you promote it.

Hyperparameter optimization

Run every trial as a task

Sweep in parallel

Stop weak trials early

Search adaptively with Optuna

Use a sweep tool you already have

Size trials for the hardware

Related