Example 15: Hyperparameter Optimization with Ray Tune¶

Manually tuning hyperparameters -- learning rate, hidden size, model type -- is tedious and error-prone. TSFast integrates with Ray Tune to automate the search. This example runs a small hyperparameter search to find the best model configuration for the Silverbox benchmark.

Prerequisites¶

This example builds on concepts from:

Example 00 -- data loading and model training basics
Example 04 -- model architectures and rnn_type

Make sure Ray Tune is installed:

uv sync --extra dev

Setup¶

Ray's uv integration requires the working directory to contain the project's pyproject.toml when Ray workers are launched. Notebook kernels start in the notebook's directory, so we move to the project root first.

In [1]:

Copied!





import os
from pathlib import Path

_root = Path.cwd()
while not (_root / "pyproject.toml").is_file() and _root != _root.parent:
    _root = _root.parent
if (_root / "pyproject.toml").is_file():
    os.chdir(_root)
import os
from pathlib import Path

_root = Path.cwd()
while not (_root / "pyproject.toml").is_file() and _root != _root.parent:
    _root = _root.parent
if (_root / "pyproject.toml").is_file():
    os.chdir(_root)

In [2]:

Copied!





import identibench as idb

from tsfast.tsdata.benchmark import create_dls_silverbox
from tsfast.tune import LearnerTrainable, ray_device, report_metrics, resume_checkpoint
from tsfast.training import RNNLearner, fun_rmse
from ray import tune
from ray.tune.schedulers import PopulationBasedTraining
import identibench as idb

from tsfast.tsdata.benchmark import create_dls_silverbox
from tsfast.tune import LearnerTrainable, ray_device, report_metrics, resume_checkpoint
from tsfast.training import RNNLearner, fun_rmse
from ray import tune
from ray.tune.schedulers import PopulationBasedTraining

/home/dom_weber/development/tsfast/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
2026-07-06 14:53:00,030	INFO util.py:155 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

2026-07-06 14:53:00,102	INFO util.py:155 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

Why Hyperparameter Optimization?¶

Model performance depends heavily on hyperparameters: learning rate, hidden size, architecture choice, and regularization strength. Finding the right combination by hand requires many experiments and careful record-keeping.

Automated approaches help:

Grid search evaluates every combination -- thorough but expensive.
Random search samples randomly and is surprisingly effective in high-dimensional spaces.
Population-based training evolves configurations during training, combining exploration with exploitation.

Ray Tune provides all of these strategies (and more) behind a unified API. TSFast provides helpers that bridge a Learner and Ray Tune:

ray_device() -- detects the GPU assigned to the current Ray worker.
report_metrics(lrn) -- patches log_epoch to report metrics and checkpoints to Ray Tune after every epoch.
resume_checkpoint(lrn) -- restores from a Ray Tune checkpoint when resuming a trial (e.g. population-based training).
LearnerTrainable -- a class-based Trainable that wraps a Learner for schedulers like Population-Based Training that need checkpoint/restore and actor reuse.

Prepare the DataLoaders¶

We use the Silverbox benchmark with a small batch size and window size to keep the example lightweight. The benchmark spec defines an initialization window that its evaluation protocol discards; using it as n_skip excludes the same initial transient from the training loss.

In [3]:

Copied!

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)
n_skip = idb.BenchmarkSilverbox_Simulation.task.init_window
dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)
n_skip = idb.BenchmarkSilverbox_Simulation.task.init_window

Define a Training Function¶

Ray Tune calls a training function once per trial, each time with a different hyperparameter configuration sampled from the search space. The DataLoaders are captured via closure.

In [4]:

Copied!





def train(config):
    lrn = RNNLearner(
        dls,
        rnn_type=config["rnn_type"],
        hidden_size=config["hidden_size"],
        n_skip=n_skip,
        metrics=[fun_rmse],
        device=ray_device(),
    )
    resume_checkpoint(lrn)
    with lrn.no_bar(), report_metrics(lrn):
        lrn.fit_flat_cos(config["n_epoch"], lr=config.get("lr"))
def train(config):
    lrn = RNNLearner(
        dls,
        rnn_type=config["rnn_type"],
        hidden_size=config["hidden_size"],
        n_skip=n_skip,
        metrics=[fun_rmse],
        device=ray_device(),
    )
    resume_checkpoint(lrn)
    with lrn.no_bar(), report_metrics(lrn):
        lrn.fit_flat_cos(config["n_epoch"], lr=config.get("lr"))

Define the Search Space¶

The search space is a plain dictionary where values are Ray Tune sampling primitives:

tune.choice -- samples uniformly from a list of discrete options. Good for categorical parameters like architecture type or layer count.
tune.loguniform -- samples uniformly on a logarithmic scale. Ideal for parameters that span orders of magnitude, such as learning rate.

Training parameters (n_epoch, lr) go in the same dict -- they are read by the training function and logged by Ray Tune.

Run the Optimization¶

Pass the training function to tune.Tuner together with the search space. num_samples=4 runs 4 independent trials, each with a different hyperparameter combination.

In [5]:

Copied!





tuner = tune.Tuner(
    tune.with_resources(train, {"cpu": 1, "gpu": 0}),
    param_space={
        "rnn_type": tune.choice(["gru", "lstm"]),
        "hidden_size": tune.choice([32, 40]),
        "n_epoch": 3,
        "lr": 3e-3,
    },
    tune_config=tune.TuneConfig(metric="valid_loss", mode="min", num_samples=4),
)

results = tuner.fit()
tuner = tune.Tuner(
    tune.with_resources(train, {"cpu": 1, "gpu": 0}),
    param_space={
        "rnn_type": tune.choice(["gru", "lstm"]),
        "hidden_size": tune.choice([32, 40]),
        "n_epoch": 3,
        "lr": 3e-3,
    },
    tune_config=tune.TuneConfig(metric="valid_loss", mode="min", num_samples=4),
)

results = tuner.fit()

Tune Status

Current time:	2026-07-06 14:53:35
Running for:	00:00:32.82
Memory:	17.3/62.6 GiB

System Info

Using FIFO scheduling algorithm.
Logical resource usage: 1.0/32 CPUs, 0/2 GPUs (0.0/1.0 accelerator_type:GeForce-RTX-4090)

Trial Status

Trial name	status	loc	hidden_size	rnn_type	iter	total time (s)	train_loss	valid_loss	fun_rmse
train_9fce9_00000	TERMINATED	130.75.144.193:850693	32	gru	3	31.0718	0.00183821	0.0010205	0.00207516
train_9fce9_00001	TERMINATED	130.75.144.193:850695	32	gru	3	30.9097	0.00163734	0.000929367	0.00201161
train_9fce9_00002	TERMINATED	130.75.144.193:850692	32	lstm	3	5.99964	0.00210793	0.00118857	0.00257935
train_9fce9_00003	TERMINATED	130.75.144.193:850694	32	gru	3	30.9489	0.00185784	0.00102332	0.00214731

(train pid=850692) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00002_2_hidden_size=32,rnn_type=lstm_2026-07-06_14-53-02/checkpoint_000000)

(train pid=850692) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00002_2_hidden_size=32,rnn_type=lstm_2026-07-06_14-53-02/checkpoint_000001)

(train pid=850692) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00002_2_hidden_size=32,rnn_type=lstm_2026-07-06_14-53-02/checkpoint_000002)

(train pid=850693) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00000_0_hidden_size=32,rnn_type=gru_2026-07-06_14-53-02/checkpoint_000000)
(train pid=850695) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00001_1_hidden_size=32,rnn_type=gru_2026-07-06_14-53-02/checkpoint_000000)

(train pid=850695) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00001_1_hidden_size=32,rnn_type=gru_2026-07-06_14-53-02/checkpoint_000001) [repeated 2x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)

(train pid=850695) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00001_1_hidden_size=32,rnn_type=gru_2026-07-06_14-53-02/checkpoint_000002) [repeated 3x across cluster]
2026-07-06 14:53:35,465	INFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/dom_weber/ray_results/train_2026-07-06_14-53-00' in 0.0025s.

2026-07-06 14:53:35,467	INFO tune.py:1039 -- Total run time: 32.83 seconds (32.81 seconds for the tuning loop).

Analyze Results¶

tuner.fit() returns a ResultGrid. You can query it for the best trial configuration, inspect per-trial metrics, or export data for further analysis.

In [6]:

Copied!





best = results.get_best_result()
print("Best config:")
for key in ["rnn_type", "hidden_size", "lr"]:
    print(f"  {key}: {best.config[key]}")
best = results.get_best_result()
print("Best config:")
for key in ["rnn_type", "hidden_size", "lr"]:
    print(f"  {key}: {best.config[key]}")

Best config:
  rnn_type: gru
  hidden_size: 32
  lr: 0.003

In [7]:

Copied!

result_df = results.get_dataframe()
print("\nAll trial results:")
result_df[["config/rnn_type", "config/hidden_size", "valid_loss"]]
result_df = results.get_dataframe()
print("\nAll trial results:")
result_df[["config/rnn_type", "config/hidden_size", "valid_loss"]]

All trial results:

Out[7]:

	config/rnn_type	config/hidden_size	valid_loss
0	gru	32	0.001020
1	gru	32	0.000929
2	lstm	32	0.001189
3	gru	32	0.001023

Using tune.loguniform for Learning Rate¶

In the first search we fixed the learning rate. A more thorough search treats lr as a tunable parameter using tune.loguniform. This samples on a logarithmic scale between the given bounds -- appropriate because the difference between 1e-4 and 1e-3 matters more than between 1e-2 and 1.1e-2.

In [8]:

Copied!





tuner_v2 = tune.Tuner(
    tune.with_resources(train, {"cpu": 1, "gpu": 0}),
    param_space={
        "rnn_type": tune.choice(["gru", "lstm"]),
        "hidden_size": tune.choice([32, 40]),
        "lr": tune.loguniform(1e-4, 1e-2),
        "n_epoch": 3,
    },
    tune_config=tune.TuneConfig(metric="valid_loss", mode="min", num_samples=4),
)

results_v2 = tuner_v2.fit()
tuner_v2 = tune.Tuner(
    tune.with_resources(train, {"cpu": 1, "gpu": 0}),
    param_space={
        "rnn_type": tune.choice(["gru", "lstm"]),
        "hidden_size": tune.choice([32, 40]),
        "lr": tune.loguniform(1e-4, 1e-2),
        "n_epoch": 3,
    },
    tune_config=tune.TuneConfig(metric="valid_loss", mode="min", num_samples=4),
)

results_v2 = tuner_v2.fit()

Tune Status

Current time:	2026-07-06 14:54:09
Running for:	00:00:34.37
Memory:	16.8/62.6 GiB

System Info

Using FIFO scheduling algorithm.
Logical resource usage: 1.0/32 CPUs, 0/2 GPUs (0.0/1.0 accelerator_type:GeForce-RTX-4090)

Trial Status

Trial name	status	loc	hidden_size	lr	rnn_type	iter	total time (s)	train_loss	valid_loss	fun_rmse
train_b3627_00000	TERMINATED	130.75.144.193:852529	40	0.000327081	gru	3	32.6527	0.0243724	0.0115071	0.0148019
train_b3627_00001	TERMINATED	130.75.144.193:852528	40	0.00247408	lstm	3	7.02782	0.00187948	0.00107462	0.00230167
train_b3627_00002	TERMINATED	130.75.144.193:852546	32	0.00481057	gru	3	30.7709	0.00235859	0.00100714	0.00212302
train_b3627_00003	TERMINATED	130.75.144.193:852560	40	0.00579735	gru	3	32.4405	0.00247291	0.00101216	0.00211003

(train pid=850693) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-00/train_9fce9_00000_0_hidden_size=32,rnn_type=gru_2026-07-06_14-53-02/checkpoint_000002)

(train pid=852528) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-35/train_b3627_00001_1_hidden_size=40,lr=0.0025,rnn_type=lstm_2026-07-06_14-53-35/checkpoint_000001) [repeated 2x across cluster]

(train pid=852546) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-35/train_b3627_00002_2_hidden_size=32,lr=0.0048,rnn_type=gru_2026-07-06_14-53-35/checkpoint_000000) [repeated 2x across cluster]

(train pid=852546) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-35/train_b3627_00002_2_hidden_size=32,lr=0.0048,rnn_type=gru_2026-07-06_14-53-35/checkpoint_000001) [repeated 3x across cluster]

(train pid=852546) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-35/train_b3627_00002_2_hidden_size=32,lr=0.0048,rnn_type=gru_2026-07-06_14-53-35/checkpoint_000002) [repeated 3x across cluster]

2026-07-06 14:54:09,858	INFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/dom_weber/ray_results/train_2026-07-06_14-53-35' in 0.0024s.

2026-07-06 14:54:09,860	INFO tune.py:1039 -- Total run time: 34.38 seconds (34.37 seconds for the tuning loop).

In [9]:

Copied!





best_v2 = results_v2.get_best_result()
print("Best config (with lr search):")
for key in ["rnn_type", "hidden_size", "lr"]:
    print(f"  {key}: {best_v2.config[key]}")
best_v2 = results_v2.get_best_result()
print("Best config (with lr search):")
for key in ["rnn_type", "hidden_size", "lr"]:
    print(f"  {key}: {best_v2.config[key]}")

Best config (with lr search):
  rnn_type: gru
  hidden_size: 32
  lr: 0.004810568824041102

Population-Based Training¶

The examples above run each trial independently. Population-Based Training (PBT) takes a different approach: it trains a population of models in parallel, periodically copying weights from the best performers and perturbing their hyperparameters. This combines the exploration of random search with the exploitation of hand-tuning.

PBT requires the class-based Trainable API so that Ray can checkpoint, restore, and mutate trials. LearnerTrainable provides this -- you supply a create_learner factory and an optional apply_config callback for actor reuse.

In [10]:

Copied!





def create_learner(config):
    return RNNLearner(
        dls,
        rnn_type="gru",
        hidden_size=32,
        n_skip=n_skip,
        metrics=[fun_rmse],
        device=ray_device(),
    )


def apply_config(lrn, config):
    for pg in lrn.opt.param_groups:
        pg["lr"] = config["lr"]
def create_learner(config):
    return RNNLearner(
        dls,
        rnn_type="gru",
        hidden_size=32,
        n_skip=n_skip,
        metrics=[fun_rmse],
        device=ray_device(),
    )


def apply_config(lrn, config):
    for pg in lrn.opt.param_groups:
        pg["lr"] = config["lr"]

In [11]:

Copied!





pbt_scheduler = PopulationBasedTraining(
    time_attr="training_iteration",
    perturbation_interval=2,
    hyperparam_mutations={"lr": tune.loguniform(1e-4, 1e-2)},
)

tuner_pbt = tune.Tuner(
    tune.with_resources(
        tune.with_parameters(
            LearnerTrainable,
            create_learner=create_learner,
            apply_config=apply_config,
        ),
        {"cpu": 1, "gpu": 0},
    ),
    param_space={"lr": 3e-3},
    tune_config=tune.TuneConfig(
        metric="valid_loss",
        mode="min",
        num_samples=4,
        scheduler=pbt_scheduler,
    ),
    run_config=tune.RunConfig(stop={"training_iteration": 6}),
)

results_pbt = tuner_pbt.fit()
pbt_scheduler = PopulationBasedTraining(
    time_attr="training_iteration",
    perturbation_interval=2,
    hyperparam_mutations={"lr": tune.loguniform(1e-4, 1e-2)},
)

tuner_pbt = tune.Tuner(
    tune.with_resources(
        tune.with_parameters(
            LearnerTrainable,
            create_learner=create_learner,
            apply_config=apply_config,
        ),
        {"cpu": 1, "gpu": 0},
    ),
    param_space={"lr": 3e-3},
    tune_config=tune.TuneConfig(
        metric="valid_loss",
        mode="min",
        num_samples=4,
        scheduler=pbt_scheduler,
    ),
    run_config=tune.RunConfig(stop={"training_iteration": 6}),
)

results_pbt = tuner_pbt.fit()

Tune Status

Current time:	2026-07-06 14:55:14
Running for:	00:01:04.46
Memory:	16.6/62.6 GiB

System Info

PopulationBasedTraining: 2 checkpoints, 0 perturbs
Logical resource usage: 1.0/32 CPUs, 0/2 GPUs (0.0/1.0 accelerator_type:GeForce-RTX-4090)

Trial Status

Trial name	status	loc	iter	total time (s)	train_loss	valid_loss	fun_rmse
LearnerTrainable_c7e19_00000	TERMINATED	130.75.144.193:854424	6	62.0051	0.00198625	0.00158017	0.00243239
LearnerTrainable_c7e19_00001	TERMINATED	130.75.144.193:854411	6	60.3968	0.00268571	0.00257701	0.0035214
LearnerTrainable_c7e19_00002	TERMINATED	130.75.144.193:854406	6	61.334	0.0020752	0.00155199	0.00234816
LearnerTrainable_c7e19_00003	TERMINATED	130.75.144.193:854437	6	61.8989	0.00216143	0.00248973	0.00336336

(train pid=852529) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/train_2026-07-06_14-53-35/train_b3627_00000_0_hidden_size=40,lr=0.0003,rnn_type=gru_2026-07-06_14-53-35/checkpoint_000002)

2026-07-06 14:54:32,651	INFO pbt.py:741 -- [pbt]: no checkpoint for trial LearnerTrainable_c7e19_00001. Skip exploit for Trial LearnerTrainable_c7e19_00003

2026-07-06 14:54:32,717	INFO pbt.py:741 -- [pbt]: no checkpoint for trial LearnerTrainable_c7e19_00001. Skip exploit for Trial LearnerTrainable_c7e19_00000

(LearnerTrainable pid=854406) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/LearnerTrainable_2026-07-06_14-54-09/LearnerTrainable_c7e19_00002_2_2026-07-06_14-54-09/checkpoint_000000)

(LearnerTrainable pid=854424) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/LearnerTrainable_2026-07-06_14-54-09/LearnerTrainable_c7e19_00000_0_2026-07-06_14-54-09/checkpoint_000000)

(LearnerTrainable pid=854411) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/LearnerTrainable_2026-07-06_14-54-09/LearnerTrainable_c7e19_00001_1_2026-07-06_14-54-09/checkpoint_000000)

(LearnerTrainable pid=854406) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/LearnerTrainable_2026-07-06_14-54-09/LearnerTrainable_c7e19_00002_2_2026-07-06_14-54-09/checkpoint_000001)

2026-07-06 14:55:14,335	INFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/dom_weber/ray_results/LearnerTrainable_2026-07-06_14-54-09' in 0.0045s.

2026-07-06 14:55:14,338	INFO tune.py:1039 -- Total run time: 64.47 seconds (64.45 seconds for the tuning loop).

In [12]:

Copied!





best_pbt = results_pbt.get_best_result()
print("Best PBT config:")
print(f"  lr: {best_pbt.config['lr']}")
print(f"  valid_loss: {best_pbt.metrics['valid_loss']:.4f}")
best_pbt = results_pbt.get_best_result()
print("Best PBT config:")
print(f"  lr: {best_pbt.config['lr']}")
print(f"  valid_loss: {best_pbt.metrics['valid_loss']:.4f}")

Best PBT config:
  lr: 0.003
  valid_loss: 0.0016

Packing Multiple Trials per GPU¶

The searches above ran on CPU. On a GPU the picture changes: a small sequence model reserves well under 1 GB of a 24 GB card and keeps the device mostly idle, so giving each trial a whole GPU wastes most of it. Ray Tune happily co-locates trials on one device ({"gpu": 1/k}) -- the only question is which k is safe and worthwhile. TSFast answers it with four functions, one per step of a measure → decide → verify → enforce workflow. They are separate on purpose:

Measure -- probe_gpu_saturation(make_learner, configs) collects the facts about your workload: per-config memory footprint (including the ~0.5 GB CUDA context every co-located process pays) and how busy a single trial keeps the device. It runs a few real training steps per config -- seconds each -- and you run it once per workload (model family × dataset).
Decide -- recommend_trials_per_gpu(probe) turns those measurements into a number: the memory ceiling k_mem, a saturation prior k_compute, and the headline rec.k = min of the two. It is pure math on the probe result -- instant, GPU-free -- so you can re-derive k under a different memory margin or max_k without measuring again.
Verify -- measure_packing_curve(make_learner, configs, ks) is the ground truth the recommendation is checked against: it launches k real training processes on one GPU and measures aggregate steps/s. The compute side of step 2 is only a prior (power draw does not see cache or memory- bandwidth contention), so before freezing k for a long tuning campaign, validate it on a few anchor configs with the curve. This is the expensive step -- minutes, not seconds -- which is exactly why it is not folded into the probe.
Enforce -- trial_resources(k) and apply_gpu_quota(share, ...) wire the frozen k into Ray Tune. Fractional GPUs in Ray are purely logical -- k trials get scheduled onto one device with no isolation -- so apply_gpu_quota, called as the first line of the trainable, caps each process's CUDA allocator to its slice. A config that grows too large then fails deterministically against its own quota instead of taking down its neighbors.

Measure and decide¶

sample_ray_space draws random probe configs from the same search space you tune over, so the probe sees the sizes a real tuning session would sample. Plain values pass through as constants.

In [13]:

Copied!





import torch

from tsfast.training.profiling import (
    probe_gpu_saturation,
    recommend_trials_per_gpu,
    sample_ray_space,
)
from tsfast.tune import apply_gpu_quota, trial_resources

search_space = {
    "rnn_type": tune.choice(["gru", "lstm"]),
    "hidden_size": tune.choice([32, 40]),
    "lr": tune.loguniform(1e-4, 1e-2),
    "n_epoch": 1,
}
import torch

from tsfast.training.profiling import (
    probe_gpu_saturation,
    recommend_trials_per_gpu,
    sample_ray_space,
)
from tsfast.tune import apply_gpu_quota, trial_resources

search_space = {
    "rnn_type": tune.choice(["gru", "lstm"]),
    "hidden_size": tune.choice([32, 40]),
    "lr": tune.loguniform(1e-4, 1e-2),
    "n_epoch": 1,
}

In [14]:

Copied!





if torch.cuda.is_available():

    def make_learner(config):
        return RNNLearner(
            dls,
            rnn_type=config["rnn_type"],
            hidden_size=config["hidden_size"],
            n_skip=n_skip,
            metrics=[fun_rmse],
        )

    probe_configs = sample_ray_space(search_space, n=8, seed=0)
    probe = probe_gpu_saturation(make_learner, probe_configs)
    rec = recommend_trials_per_gpu(probe, mem_margin=0.9, max_k=8)
    print(f"k={rec.k}  (memory ceiling {rec.k_mem}, saturation prior {rec.k_compute})")
if torch.cuda.is_available():

    def make_learner(config):
        return RNNLearner(
            dls,
            rnn_type=config["rnn_type"],
            hidden_size=config["hidden_size"],
            n_skip=n_skip,
            metrics=[fun_rmse],
        )

    probe_configs = sample_ray_space(search_space, n=8, seed=0)
    probe = probe_gpu_saturation(make_learner, probe_configs)
    rec = recommend_trials_per_gpu(probe, mem_margin=0.9, max_k=8)
    print(f"k={rec.k}  (memory ceiling {rec.k_mem}, saturation prior {rec.k_compute})")

(LearnerTrainable pid=854424) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/dom_weber/ray_results/LearnerTrainable_2026-07-06_14-54-09/LearnerTrainable_c7e19_00000_0_2026-07-06_14-54-09/checkpoint_000001)

=== GPU Saturation Probe ===
Device: cuda:0 (23.5 GB) | 8 configs
Footprint: median 0.54 GB | max 0.54 GB | ratio 1.00 | context 0.46 GB
Busy fraction: 0.04 | power fraction: 0.11
k=8  (memory ceiling 38, saturation prior 9)

Verify¶

Before a multi-day tuning campaign, check the recommendation against the measured packing curve on a few anchor configs. measure_packing_curve spawns fresh worker processes that must import make_learner, so it runs from a script (module-level factory), not from a notebook cell:

curve = measure_packing_curve(make_learner, probe_configs[:3], ks=(1, 2, 4))
k = min(rec.k, curve.knee)   # freeze this number for the campaign

On 2× RTX 4090 a small GRU workload measured 139/212/375 aggregate steps/s at k=1/2/4 -- co-locating four trials nearly tripled tuning throughput.

Enforce¶

trial_resources(k) makes Ray schedule k trials per GPU; apply_gpu_quota at the top of the trainable makes each trial's memory slice binding. The context overhead measured by the probe is passed along so the k allocator quotas plus the k CUDA contexts fit the device.

In [15]:

Copied!





if torch.cuda.is_available():

    def train_packed(config):
        apply_gpu_quota(1 / rec.k, context_bytes=rec.context_overhead_bytes)
        lrn = RNNLearner(
            dls,
            rnn_type=config["rnn_type"],
            hidden_size=config["hidden_size"],
            n_skip=n_skip,
            metrics=[fun_rmse],
            device=ray_device(),
        )
        with lrn.no_bar(), report_metrics(lrn, checkpoint_every=None):
            lrn.fit_flat_cos(config["n_epoch"], lr=config.get("lr"))

    tuner_packed = tune.Tuner(
        tune.with_resources(train_packed, trial_resources(rec.k)),
        param_space=search_space,
        tune_config=tune.TuneConfig(metric="valid_loss", mode="min", num_samples=4),
    )
    results_packed = tuner_packed.fit()
    print(f"Best packed-run loss: {results_packed.get_best_result().metrics['valid_loss']:.4f}")
if torch.cuda.is_available():

    def train_packed(config):
        apply_gpu_quota(1 / rec.k, context_bytes=rec.context_overhead_bytes)
        lrn = RNNLearner(
            dls,
            rnn_type=config["rnn_type"],
            hidden_size=config["hidden_size"],
            n_skip=n_skip,
            metrics=[fun_rmse],
            device=ray_device(),
        )
        with lrn.no_bar(), report_metrics(lrn, checkpoint_every=None):
            lrn.fit_flat_cos(config["n_epoch"], lr=config.get("lr"))

    tuner_packed = tune.Tuner(
        tune.with_resources(train_packed, trial_resources(rec.k)),
        param_space=search_space,
        tune_config=tune.TuneConfig(metric="valid_loss", mode="min", num_samples=4),
    )
    results_packed = tuner_packed.fit()
    print(f"Best packed-run loss: {results_packed.get_best_result().metrics['valid_loss']:.4f}")

Tune Status

Current time:	2026-07-06 14:55:19
Running for:	00:00:03.23
Memory:	19.7/62.6 GiB

System Info

Using FIFO scheduling algorithm.
Logical resource usage: 1.0/32 CPUs, 0.125/2 GPUs (0.0/1.0 accelerator_type:GeForce-RTX-4090)

Trial Status

Trial name	status	loc	hidden_size	lr	rnn_type	iter	total time (s)	train_loss	valid_loss	fun_rmse
train_packed_ef701_00000	TERMINATED	130.75.144.193:857733	40	0.00369007	gru	1	1.39104	0.0165099	0.00125073	0.00234595
train_packed_ef701_00001	TERMINATED	130.75.144.193:857761	32	0.000338333	lstm	1	1.51785	0.0405544	0.0357397	0.0446991
train_packed_ef701_00002	TERMINATED	130.75.144.193:857762	40	0.000148443	gru	1	1.52174	0.0512869	0.0385149	0.048186
train_packed_ef701_00003	TERMINATED	130.75.144.193:857776	40	0.0013205	gru	1	1.46093	0.0162647	0.00160978	0.00265622

2026-07-06 14:55:19,473	INFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/dom_weber/ray_results/train_packed_2026-07-06_14-55-16' in 0.0024s.

2026-07-06 14:55:19,475	INFO tune.py:1039 -- Total run time: 3.24 seconds (3.22 seconds for the tuning loop).

Best packed-run loss: 0.0013

Key Takeaways¶

ray_device() detects the GPU assigned to the current Ray worker -- pass it as device when constructing your Learner.
report_metrics(lrn) patches log_epoch to report metrics and checkpoints to Ray Tune after every epoch.
resume_checkpoint(lrn) restores from a Ray Tune checkpoint when resuming a trial (needed for population-based training).
LearnerTrainable wraps a Learner as a class-based Trainable for schedulers like PBT that need checkpoint/restore and actor reuse. Supply a create_learner factory and an optional apply_config callback via tune.with_parameters.
tune.choice for categorical parameters (architecture, layer count); tune.loguniform for continuous parameters on a log scale (learning rate). No custom samplers needed.
tune.Tuner gives you full control over scheduling, stopping criteria, and resource allocation.
GPU packing follows measure → decide → verify → enforce: probe_gpu_saturation measures footprints and saturation, recommend_trials_per_gpu derives k from them, measure_packing_curve validates k against ground-truth aggregate throughput, and trial_resources + apply_gpu_quota wire the frozen k into Ray Tune.
Start small -- few trials, few epochs -- to validate the pipeline before scaling up.